black reformat

kube_hunter/__main__.py

@@ -8,7 +8,11 @@ from kube_hunter.conf import Config, set_config
 from kube_hunter.conf.parser import parse_args
 from kube_hunter.conf.logging import setup_logger
 
-args = parse_args()
+from kube_hunter.plugins import initialize_plugin_manager
+
+pm = initialize_plugin_manager()
+# Using a plugin hook for adding arguments before parsing
+args = parse_args(add_args_hook=pm.hook.parser_add_arguments)
 config = Config(
     active=args.active,
     cidr=args.cidr,
@@ -24,6 +28,9 @@ config = Config(
 setup_logger(args.log)
 set_config(config)
 
+# Running all other registered plugins before execution
+pm.hook.load_plugin(args=args)
+
 from kube_hunter.core.events import handler
 from kube_hunter.core.events.types import HuntFinished, HuntStarted
 from kube_hunter.modules.discovery.hosts import RunningAsPodEvent, HostScanEvent

kube_hunter/conf/parser.py

@@ -1,9 +1,13 @@
 from argparse import ArgumentParser
+from kube_hunter.plugins import hookimpl
 
 
-def parse_args():
-    parser = ArgumentParser(description="kube-hunter - hunt for security weaknesses in Kubernetes clusters")
-
+@hookimpl
+def parser_add_arguments(parser):
+    """
+    This is the default hook implementation for parse_add_argument
+    Contains initialization for all default arguments
+    """
     parser.add_argument(
         "--list", action="store_true", help="Displays all tests in kubehunter (add --active flag to see active tests)",
     )
@@ -59,6 +63,18 @@ def parse_args():
 
     parser.add_argument("--network-timeout", type=float, default=5.0, help="network operations timeout")
 
+
+def parse_args(add_args_hook):
+    """
+    Function handles all argument parsing
+
+    @param add_arguments: hook for adding arguments to it's given ArgumentParser parameter
+    @return: parsed arguments dict
+    """
+    parser = ArgumentParser(description="kube-hunter - hunt for security weaknesses in Kubernetes clusters")
+    # adding all arguments to the parser
+    add_args_hook(parser=parser)
+
     args = parser.parse_args()
     if args.cidr:
         args.cidr = args.cidr.replace(" ", "").split(",")

kube_hunter/plugins/__init__.py

@@ -0,0 +1,23 @@
+import pluggy
+
+from kube_hunter.plugins import hookspecs
+
+hookimpl = pluggy.HookimplMarker("kube-hunter")
+
+
+def initialize_plugin_manager():
+    """
+    Initializes and loads all default and setup implementations for registered plugins
+
+    @return: initialized plugin manager
+    """
+    pm = pluggy.PluginManager("kube-hunter")
+    pm.add_hookspecs(hookspecs)
+    pm.load_setuptools_entrypoints("kube_hunter")
+
+    # default registration of builtin implemented plugins
+    from kube_hunter.conf import parser
+
+    pm.register(parser)
+
+    return pm

kube_hunter/plugins/hookspecs.py

@@ -0,0 +1,24 @@
+import pluggy
+from argparse import ArgumentParser
+
+hookspec = pluggy.HookspecMarker("kube-hunter")
+
+
+@hookspec
+def parser_add_arguments(parser: ArgumentParser):
+    """Add arguments to the ArgumentParser.
+
+    If a plugin requires an aditional argument, it should implement this hook
+    and add the argument to the Argument Parser
+
+    @param parser: an ArgumentParser, calls parser.add_argument on it
+    """
+
+
+@hookspec
+def load_plugin(args):
+    """Plugins that wish to execute code after the argument parsing
+    should implement this hook.
+
+    @param args: all parsed arguments passed to kube-hunter
+    """

requirements-dev.txt

@@ -14,3 +14,4 @@ black
 pre-commit
 flake8-bugbear
 flake8-mypy
+pluggy

setup.cfg

@@ -38,6 +38,7 @@ install_requires =
     future
     packaging
     dataclasses
+    pluggy
 setup_requires =
     setuptools>=30.3.0
     setuptools_scm

tests/plugins/test_hooks.py

@@ -0,0 +1,13 @@
+from kube_hunter.plugins import hookimpl
+
+return_string = "return_string"
+
+
+@hookimpl
+def parser_add_arguments(parser):
+    return return_string
+
+
+@hookimpl
+def load_plugin(args):
+    return return_string

tests/plugins/test_plugins_hooks.py

@@ -0,0 +1,17 @@
+from argparse import ArgumentParser
+from tests.plugins import test_hooks
+from kube_hunter.plugins import initialize_plugin_manager
+
+
+def test_all_plugin_hooks():
+    pm = initialize_plugin_manager()
+    pm.register(test_hooks)
+
+    # Testing parser_add_arguments
+    parser = ArgumentParser("Test Argument Parser")
+    results = pm.hook.parser_add_arguments(parser=parser)
+    assert test_hooks.return_string in results
+
+    # Testing load_plugin
+    results = pm.hook.load_plugin(args=[])
+    assert test_hooks.return_string in results