From 3b6f067f9352d2842157b4569eca2640eb24106c Mon Sep 17 00:00:00 2001 From: Miauwkeru Date: Thu, 18 Jul 2024 12:40:01 +0000 Subject: [PATCH] Add optional boolean value to all store_true instances (#179) --- acquire/acquire.py | 16 +++++++--------- acquire/utils.py | 16 +++++++++------- tests/test_acquire_command.py | 25 ++++++++++++++++++++----- 3 files changed, 36 insertions(+), 21 deletions(-) diff --git a/acquire/acquire.py b/acquire/acquire.py index 77b5db6a..b6ce64bb 100644 --- a/acquire/acquire.py +++ b/acquire/acquire.py @@ -215,7 +215,7 @@ def wrapper(module_cls: type[Module]) -> type[Module]: desc = module_cls.DESC or name kwargs["help"] = f"acquire {desc}" - kwargs["action"] = "store_true" + kwargs["action"] = argparse.BooleanOptionalAction kwargs["dest"] = name.lower() module_cls.__modname__ = name @@ -661,15 +661,13 @@ def recyclebin_filter(path: fsutil.TargetPath) -> bool: @register_module("--recyclebin") @module_arg( "--large-files", - action="store_true", + action=argparse.BooleanOptionalAction, help="Collect files larger than 10MB in the Recycle Bin", - default=False, ) @module_arg( - "--no-data-files", - action="store_true", - help="Skip collection of data files in the Recycle Bin", - default=False, + "--data-files", + action=argparse.BooleanOptionalAction, + help="Collect the data files in the Recycle Bin", ) class RecycleBin(Module): DESC = "recycle bin metadata and data files" @@ -683,7 +681,7 @@ def _run(cls, target: Target, cli_args: argparse.Namespace, collector: Collector patterns = ["$Recycle.bin/*/$I*", "Recycler/*/INFO2", "Recycled/INFO2"] - if not cli_args.no_data_files: + if cli_args.data_files is None or cli_args.data_files: patterns.extend(["$Recycle.Bin/$R*", "$Recycle.Bin/*/$R*", "RECYCLE*/D*"]) with collector.file_filter(large_files_filter): @@ -1328,7 +1326,7 @@ class Home(Module): @register_module("--ssh") -@module_arg("--private-keys", action="store_true", help="Add any private keys", default=False) +@module_arg("--private-keys", action=argparse.BooleanOptionalAction, help="Add any private keys") class SSH(Module): SPEC = [ ("glob", ".ssh/*", from_user_home), diff --git a/acquire/utils.py b/acquire/utils.py index 034eb807..e6d0f5ff 100644 --- a/acquire/utils.py +++ b/acquire/utils.py @@ -81,12 +81,12 @@ def create_argument_parser(profiles: dict, volatile: dict, modules: dict) -> arg ) parser.add_argument( "--compress", - action="store_true", + action=argparse.BooleanOptionalAction, help="compress output (if supported by the output type)", ) parser.add_argument( "--encrypt", - action="store_true", + action=argparse.BooleanOptionalAction, help="encrypt output (if supported by the output type)", ) parser.add_argument( @@ -119,19 +119,21 @@ def create_argument_parser(profiles: dict, volatile: dict, modules: dict) -> arg parser.add_argument("--child", help="only collect specific child") parser.add_argument( "--children", - action="store_true", + action=argparse.BooleanOptionalAction, help="collect all children in addition to main target", ) - parser.add_argument("--skip-parent", action="store_true", help="skip parent collection (when using --children)") + parser.add_argument( + "--skip-parent", action=argparse.BooleanOptionalAction, help="skip parent collection (when using --children)" + ) parser.add_argument( "--force-fallback", - action="store_true", + action=argparse.BooleanOptionalAction, help="force filesystem access directly through OS level. Only supported with target 'local'", ) parser.add_argument( "--fallback", - action="store_true", + action=argparse.BooleanOptionalAction, help=( "fallback to OS level filesystem access if filesystem type is not supported. " "Only supported with target 'local'" @@ -141,7 +143,7 @@ def create_argument_parser(profiles: dict, volatile: dict, modules: dict) -> arg parser.add_argument( "-u", "--auto-upload", - action="store_true", + action=argparse.BooleanOptionalAction, help="upload result files after collection", ) parser.add_argument( diff --git a/tests/test_acquire_command.py b/tests/test_acquire_command.py index 73a20bcf..b1d2cb6e 100644 --- a/tests/test_acquire_command.py +++ b/tests/test_acquire_command.py @@ -1,11 +1,9 @@ from argparse import Namespace -from typing import List from unittest.mock import patch import pytest from acquire.acquire import ( - CONFIG, MODULES, PROFILES, VOLATILE, @@ -15,10 +13,11 @@ @pytest.fixture -def acquire_parser_args(config: List, argument_list: List) -> Namespace: - CONFIG["arguments"] = config +def acquire_parser_args(config: list[str], argument_list: list[str]) -> Namespace: + config_dict = {} + config_dict["arguments"] = config with patch("argparse._sys.argv", [""] + argument_list): - return parse_acquire_args(create_argument_parser(PROFILES, VOLATILE, MODULES), config=CONFIG)[0] + return parse_acquire_args(create_argument_parser(PROFILES, VOLATILE, MODULES), config=config_dict)[0] @pytest.mark.parametrize("config, argument_list", [([], [])]) @@ -39,3 +38,19 @@ def test_config_default_argument_override(acquire_parser_args): @pytest.mark.parametrize("config, argument_list", [([], ["target1", "target2"])]) def test_local_target_fallbactargets(acquire_parser_args): assert acquire_parser_args.targets == ["target1", "target2"] + + +@pytest.mark.parametrize( + "config, argument_list, arg_to_test, expected_value", + [ + (["--etc"], ["--no-etc"], "etc", False), + (["--no-etc"], ["--etc"], "etc", True), + (["--encrypt"], ["--no-encrypt"], "encrypt", False), + (["--no-encrypt"], ["--encrypt"], "encrypt", True), + (["--encrypt", "--ssh"], ["--no-ssh"], "ssh", False), + (["--private-keys"], ["--no-private-keys"], "private_keys", False), + (["--no-private-keys"], ["--private-keys"], "private_keys", True), + ], +) +def test_overwrites_optionals(acquire_parser_args: Namespace, arg_to_test: str, expected_value: bool): + assert getattr(acquire_parser_args, arg_to_test) is expected_value