diff --git a/dissect/target/plugins/filesystem/ntfs/mft.py b/dissect/target/plugins/filesystem/ntfs/mft.py index e61a9c3ef..417d7d911 100644 --- a/dissect/target/plugins/filesystem/ntfs/mft.py +++ b/dissect/target/plugins/filesystem/ntfs/mft.py @@ -138,12 +138,18 @@ def check_compatible(self) -> None: FilesystemFilenameCompactRecord, ] ) - @arg("--compact", action="store_true", help="compacts the MFT entry timestamps into a single record") + @arg( + "--compact", + group="fmt", + action="store_true", + help="compacts the MFT entry timestamps into a single record", + ) @arg("--fs", type=int, default=None, help="optional filesystem index, zero indexed") @arg("--start", type=int, default=0, help="the first MFT segment number") @arg("--end", type=int, default=-1, help="the last MFT segment number") @arg( "--macb", + group="fmt", action="store_true", help="compacts the MFT entry timestamps into aggregated records with MACB bitfield", ) @@ -171,9 +177,7 @@ def noaggr(records: list[Record]) -> Iterator[Record]: aggr = noaggr - if compact and macb: - raise ValueError("--macb and --compact are mutually exclusive") - elif compact: + if compact: record_formatter = compacted_formatter elif macb: aggr = macb_aggr diff --git a/dissect/target/tools/utils.py b/dissect/target/tools/utils.py index d69317cca..a27ee8676 100644 --- a/dissect/target/tools/utils.py +++ b/dissect/target/tools/utils.py @@ -95,12 +95,24 @@ def generate_argparse_for_unbound_method( parser = argparse.ArgumentParser(description=desc, formatter_class=help_formatter, conflict_handler="resolve") fargs = getattr(method, "__args__", []) + groups = {} + default_group_options = {"required": False} for args, kwargs in fargs: - parser.add_argument(*args, **kwargs) + if "group" in kwargs: + group_name = kwargs.pop("group") + options = kwargs.pop("group_options") if "group_options" in kwargs else default_group_options + if group_name not in groups: + group = parser.add_mutually_exclusive_group(**options) + groups[group_name] = group + else: + group = groups[group_name] + + group.add_argument(*args, **kwargs) + else: + parser.add_argument(*args, **kwargs) usage = parser.format_usage() offset = usage.find(parser.prog) + len(parser.prog) - func_name = method.__name__ usage_tmpl = usage_tmpl or "{prog} {usage}" parser.usage = usage_tmpl.format(prog=parser.prog, name=func_name, usage=usage[offset:]) diff --git a/tests/tools/test_utils.py b/tests/tools/test_utils.py index d7066dbc0..db5975014 100644 --- a/tests/tools/test_utils.py +++ b/tests/tools/test_utils.py @@ -8,6 +8,7 @@ from dissect.target.plugin import arg, find_plugin_functions from dissect.target.tools.utils import ( args_to_uri, + generate_argparse_for_unbound_method, get_target_attribute, persist_execution_report, ) @@ -77,3 +78,17 @@ def test_plugin_name_confusion_regression(target_unix_users, pattern, expected_f get_target_attribute(target_unix_users, plugins[0]) assert expected_function in str(exc_info.value) + + +def test_plugin_mutual_exclusive_arguments(): + fargs = [ + (("--aa",), {"group": "aa"}), + (("--bb",), {"group": "aa"}), + (("--cc",), {"group": "bb"}), + (("--dd",), {"group": "bb"}), + ] + method = test_plugin_mutual_exclusive_arguments + setattr(method, "__args__", fargs) + with patch("inspect.isfunction", return_value=True): + parser = generate_argparse_for_unbound_method(method) + assert len(parser._mutually_exclusive_groups) == 2