From 93ad35d083b7bb023bdb030fbe73092e28b9aae7 Mon Sep 17 00:00:00 2001
From: benyissa <benyissa@ostorlab.dev>
Date: Tue, 21 Nov 2023 09:55:26 +0100
Subject: [PATCH] refactor

---
 agent/exploits/__init__.py | 65 ++------------------------------------
 agent/importer.py          | 53 +++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 62 deletions(-)
 create mode 100644 agent/importer.py

diff --git a/agent/exploits/__init__.py b/agent/exploits/__init__.py
index 4fb1723c..88f8806c 100644
--- a/agent/exploits/__init__.py
+++ b/agent/exploits/__init__.py
@@ -1,63 +1,4 @@
-"""Exploits importer"""
-import logging
-import sys
-import traceback
-import os
-import pkgutil
-from typing import Callable
+"""Importer module."""
+from agent import importer
 
-logger = logging.getLogger(__name__)
-
-
-def _importer_for(path: str, prefix: str) -> Callable[[], None]:
-    """
-    Creates a function for importing all Python files in a specified folder and adding them as attributes
-    to a given module.
-
-    Args:
-        path (str): the path to the folder containing the Python files to import
-        prefix (str): the prefix to use when adding the imported files as attributes to the module
-
-    Returns:
-        A function that can be called to perform the import and attribute assignment.
-        The function takes two optional arguments:
-        - file_path (str): the path to the folder containing the Python files to import (default is `path`)
-        - stop_on_error (bool): whether to stop execution if an error occurs during import (default is `True`)
-    """
-
-    def _import_all_files(file_path: str = path, stop_on_error: bool = True) -> None:
-        """
-        Imports all Python files in a specified folder and adds them as attributes to a given module.
-
-        Args:
-            file_path (str): the path to the folder containing the Python files to import
-                            (default is the value passed to `_importer_for`)
-            stop_on_error (bool): whether to stop execution if an error occurs during import (default is `True`)
-        """
-        folder = os.path.dirname(file_path)
-        module = sys.modules[prefix]
-
-        for importer, name, _ in pkgutil.iter_modules([folder]):
-            abs_name = prefix + "." + name
-
-            if abs_name in sys.modules:
-                continue
-
-            loader = importer.find_module(abs_name)  # type: ignore[call-arg]
-            sub_mod = None
-            try:
-                if loader is not None:
-                    sub_mod = loader.load_module(abs_name)
-            except ImportError as e:
-                if stop_on_error:
-                    raise
-                # Print the full traceback for debugging purposes.
-                traceback.print_exc()
-                logger.warning("Cannot load REbus plugin [%s]. Root cause: %s", name, e)
-            else:
-                setattr(module, name, sub_mod)
-
-    return _import_all_files
-
-
-import_all = _importer_for(__file__, __name__)
+import_all = importer.importer_for(__file__, __name__)
diff --git a/agent/importer.py b/agent/importer.py
new file mode 100644
index 00000000..92cd9810
--- /dev/null
+++ b/agent/importer.py
@@ -0,0 +1,53 @@
+"""Dynamic importer to enable registry pattern."""
+import logging
+import sys
+import traceback
+import os
+import pkgutil
+from typing import Callable
+
+logger = logging.getLogger(__name__)
+
+
+def importer_for(path: str, prefix: str) -> Callable[[], None]:
+    """
+    Creates a function for importing all Python files in a specified folder and adding them as attributes
+    to a given module.
+    Args:
+        path (str): the path to the folder containing the Python files to import
+        prefix (str): the prefix to use when adding the imported files as attributes to the module
+    Returns:
+        A function that can be called to perform the import and attribute assignment. The function takes two
+        optional arguments:
+        - file_path (str): the path to the folder containing the Python files to import (default is `path`)
+        - stop_on_error (bool): whether to stop execution if an error occurs during import (default is `True`)
+    """
+
+    def import_all(path: str = path, stop_on_error: bool = True) -> None:
+        """
+        Imports all Python files in a specified folder and adds them as attributes to a given module.
+        Args:
+            file_path (str): the path to the folder containing the Python files to import
+             (default is the value passed to `importer_for`)
+            stop_on_error (bool): whether to stop execution if an error occurs during import (default is `True`)
+        """
+        folder = os.path.dirname(path)
+        module = sys.modules[prefix]
+        for importer, name, _ in pkgutil.iter_modules([folder]):
+            absname = prefix + "." + name
+            if absname in sys.modules:
+                continue
+            loader = importer.find_module(absname)  # type: ignore[call-arg]
+            try:
+                if loader is not None:
+                    submod = loader.load_module(absname)
+            except ImportError as e:
+                if stop_on_error:
+                    raise
+                # This is for debugging to print the full trace and pinpoint to source of the exception.
+                traceback.print_exc()
+                logger.warning("Cannot load REbus plugin [%s]. Root cause: %s", name, e)
+            else:
+                setattr(module, name, submod)
+
+    return import_all