Add optional boolean value to all store_true instances (#179)

fox-it · Jul 18, 2024 · 3b6f067 · 3b6f067
1 parent 783b1ec
commit 3b6f067
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 21 deletions.
diff --git 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
@@ -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
@@ -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