Improve parsing of plugin arguments in target-query

dissect/target/plugins/filesystem/walkfs.py

@@ -50,7 +50,7 @@ def check_compatible(self) -> None:
             raise UnsupportedPluginError("No filesystems to walk")
 
     @export(record=FilesystemRecord)
-    @arg("--walkfs-path", default="/", help="path to recursively walk")
+    @arg("-p", "--walkfs-path", default="/", help="path to recursively walk")
     @arg("--capability", action="store_true", help="output capability records")
     @arg("--mimetype", action="store_true", help="enable mimetype lookup of files")
     def walkfs(

dissect/target/tools/query.py

@@ -71,8 +71,10 @@ def main() -> int:
         fromfile_prefix_chars="@",
         formatter_class=help_formatter,
         add_help=False,
+        conflict_handler="resolve",
     )
     parser.add_argument("targets", metavar="TARGETS", nargs="*", help="Targets to load")
+    parser.add_argument("-h", "--help", action="store_true", help="show this help message and exit")
     parser.add_argument("--direct", action="store_true", help="treat TARGETS as paths to pass to plugins directly")
 
     configure_plugin_arguments(parser)
@@ -113,11 +115,32 @@ def main() -> int:
 
     args, rest = parser.parse_known_args()
 
-    # Show help for target-query
-    if not args.function and ("-h" in rest or "--help" in rest):
+    # Show help for target-query (when --help is specified without a function)
+    if not args.function and args.help:
         parser.print_help()
         return 0
 
+    # Dynamically add plugin arguments for the specified function(s)
+    for func_desc in find_and_filter_plugins(
+        functions=args.function or "",
+        target=None,
+        excluded_func_paths=args.excluded_functions,
+    ):
+        plugin_args = func_desc.args
+        for opts, kwargs in plugin_args:
+            # Skip 'group' if it exists, as this is a special argument used in dissect.
+            kwargs = {k: v for k, v in kwargs.items() if k != "group"}
+            parser.add_argument(*opts, **kwargs)
+
+    # Re-parse arguments now that requested plugin arguments have been added.
+    # Unknown args will automatically be shown to the user here to catch any misspelled or invalid arguments early on.
+    args = parser.parse_args()
+
+    # Propagate --help argument to plugin to show the help message of the plugin.
+    rest = []
+    if args.help:
+        rest.append("--help")
+
     process_generic_arguments(parser, args)
 
     if args.no_cache:

tests/tools/test_query.py

@@ -4,7 +4,7 @@
 import json
 import re
 from typing import TYPE_CHECKING, Any
-from unittest.mock import patch
+from unittest.mock import MagicMock, PropertyMock, patch
 
 import pytest
 
@@ -220,7 +220,21 @@ def mock_execute_function(
     return (func.output, func.name)
 
 
-def test_filtered_functions(monkeypatch: pytest.MonkeyPatch) -> None:
+@patch("dissect.target.plugin.PLUGINS", new_callable=PluginRegistry)
+def test_filtered_functions(mock_plugins: PluginRegistry, monkeypatch: pytest.MonkeyPatch) -> None:
+    class MockQueryArgsPlugin(Plugin):
+        def check_compatible(self) -> None:
+            pass
+
+        @export(output="none")
+        @arg("-j", "--json", action="store_true")
+        @arg("--compact", action="store_true")
+        def mock_function(self, json: bool, compact: bool) -> None:
+            pass
+
+    def mock_query_plugin_args() -> list[tuple[list[str], dict[str, Any]]]:
+        return getattr(MockQueryArgsPlugin(MagicMock()).mock_function, "__args__", [])
+
     with monkeypatch.context() as m:
         m.setattr(
             "sys.argv",
@@ -235,6 +249,12 @@ def test_filtered_functions(monkeypatch: pytest.MonkeyPatch) -> None:
         )
 
         with (
+            patch.object(
+                FunctionDescriptor,
+                "args",
+                new_callable=PropertyMock,
+                side_effect=mock_query_plugin_args,
+            ),
             patch(
                 "dissect.target.tools.utils.cli.find_functions",
                 autospec=True,
@@ -441,6 +461,9 @@ def mock_function(self, json: bool, compact: bool) -> None:
             assert json is True
             assert compact is True
 
+    def mock_query_plugin_args() -> list[tuple[list[str], dict[str, Any]]]:
+        return getattr(MockPlugin(MagicMock()).mock_function, "__args__", [])
+
     with monkeypatch.context() as m:
         m.setattr(
             "sys.argv",
@@ -454,6 +477,12 @@ def mock_function(self, json: bool, compact: bool) -> None:
         )
 
         with (
+            patch.object(
+                FunctionDescriptor,
+                "args",
+                new_callable=PropertyMock,
+                side_effect=mock_query_plugin_args,
+            ),
             patch(
                 "dissect.target.tools.utils.cli.find_functions",
                 autospec=True,
@@ -511,3 +540,53 @@ def test_direct_mode(capsys: pytest.CaptureFixture, monkeypatch: pytest.MonkeyPa
         target_query()
         out, _ = capsys.readouterr()
     assert len(out.splitlines()) == 3
+
+
+def test_plugin_argument_handling(capsys: pytest.CaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Test argument handling of target-query for a plugin with arguments."""
+    with monkeypatch.context() as m:
+        m.setattr(
+            "sys.argv",
+            [
+                "target-query",
+                "-q",
+                "-s",  # output as text records
+                "-f",
+                "walkfs",
+                "--walkfs-path",
+                "/path/to/symlink",
+                str(absolute_path("_data/filesystems/symlink_disk.ext4")),
+            ],
+        )
+
+        target_query()
+
+    out, _ = capsys.readouterr()
+    assert "ino=14" in out
+    assert "ino=15" in out
+    assert len(out.splitlines()) == 2
+
+
+def test_plugin_argument_unknown(capsys: pytest.CaptureFixture, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Test that unknown plugin arguments are caught and logged."""
+    with monkeypatch.context() as m:
+        m.setattr(
+            "sys.argv",
+            [
+                "target-query",
+                "-q",
+                "-s",  # output as text records
+                "-f",
+                "walkfs,yara",
+                "--some-unknown-arg=/usr/local/bin",
+                "--rules",
+                str(absolute_path("tests/_data/plugins/filesystem/yara/rule-dir/rule.yar")),
+                str(absolute_path("_data/filesystems/symlink_disk.ext4")),
+            ],
+        )
+
+        with pytest.raises(SystemExit):
+            target_query()
+
+    _, err = capsys.readouterr()
+    assert "error: unrecognized arguments: --some-unknown-arg" in err