Add support for plugin directories

fox-it · Jan 29, 2025 · 615b2bc · 615b2bc
1 parent de93169
commit 615b2bc
Show file tree

Hide file tree

Showing 2 changed files with 98 additions and 0 deletions.
diff --git a/dissect/target/plugin.py b/dissect/target/plugin.py
@@ -485,6 +485,12 @@ def register(plugincls: type[Plugin]) -> None:
     exports = []
     functions = []
     module_path = _module_path(plugincls)
+
+    # This enables plugin directories, e.g.:
+    # <plugin>/_plugin.py
+    # <plugin>/helpers.py
+    module_path = module_path.removesuffix("._plugin")
+
     module_key = f"{module_path}.{plugincls.__qualname__}"
 
     if not issubclass(plugincls, ChildTargetPlugin):

diff --git a/tests/test_plugin.py b/tests/test_plugin.py
@@ -1,4 +1,5 @@
 import os
+import textwrap
 from functools import reduce
 from pathlib import Path
 from typing import Iterator, Optional
@@ -31,6 +32,7 @@
     find_plugin_functions,
     get_external_module_paths,
     load,
+    load_modules_from_paths,
     plugins,
 )
 from dissect.target.plugins.os.default._os import DefaultPlugin
@@ -77,6 +79,96 @@ def test_load_paths_with_env() -> None:
         assert get_external_module_paths([Path(""), Path("")]) == [Path("")]
 
 
+@patch("dissect.target.plugin.PLUGINS", new_callable=PluginRegistry)
+def test_plugin_directory(mock_plugins: PluginRegistry, tmp_path: Path) -> None:
+    code = """
+        from dissect.target.plugin import Plugin, export
+
+        class MyPlugin(Plugin):
+            __namespace__ = {!r}
+
+            @export
+            def my_function(self):
+                return "My function"
+    """
+
+    (tmp_path / "myplugin").mkdir()
+    (tmp_path / "myplugin" / "__init__.py").write_text("")
+    (tmp_path / "myplugin" / "_plugin.py").write_text(textwrap.dedent(code.format(None)))
+
+    (tmp_path / "mypluginns").mkdir()
+    (tmp_path / "mypluginns" / "__init__.py").write_text("")
+    (tmp_path / "mypluginns" / "_plugin.py").write_text(textwrap.dedent(code.format("myns")))
+
+    load_modules_from_paths([tmp_path])
+
+    assert mock_plugins.__functions__.__regular__ == {
+        "my_function": {
+            "myplugin.MyPlugin": FunctionDescriptor(
+                name="my_function",
+                namespace=None,
+                path="myplugin.my_function",
+                exported=True,
+                internal=False,
+                findable=True,
+                output="default",
+                method_name="my_function",
+                module="myplugin._plugin",
+                qualname="MyPlugin",
+            )
+        },
+        "myns": {
+            "mypluginns.MyPlugin": FunctionDescriptor(
+                name="myns",
+                namespace="myns",
+                path="mypluginns",
+                exported=True,
+                internal=False,
+                findable=True,
+                output=None,
+                method_name="__call__",
+                module="mypluginns._plugin",
+                qualname="MyPlugin",
+            )
+        },
+        "myns.my_function": {
+            "mypluginns.MyPlugin": FunctionDescriptor(
+                name="myns.my_function",
+                namespace="myns",
+                path="mypluginns.my_function",
+                exported=True,
+                internal=False,
+                findable=True,
+                output="default",
+                method_name="my_function",
+                module="mypluginns._plugin",
+                qualname="MyPlugin",
+            )
+        },
+    }
+
+    assert mock_plugins.__plugins__.__regular__ == {
+        "myplugin.MyPlugin": PluginDescriptor(
+            module="myplugin._plugin",
+            qualname="MyPlugin",
+            namespace=None,
+            path="myplugin",
+            findable=True,
+            functions=["my_function"],
+            exports=["my_function"],
+        ),
+        "mypluginns.MyPlugin": PluginDescriptor(
+            module="mypluginns._plugin",
+            qualname="MyPlugin",
+            namespace="myns",
+            path="mypluginns",
+            findable=True,
+            functions=["my_function", "__call__"],
+            exports=["my_function", "__call__"],
+        ),
+    }
+
+
 class MockOSWarpPlugin(OSPlugin):
     __exports__ = ["f6"]  # OS exports f6
     __register__ = False