Patched patchwork/app.py

Author: patched.codes[bot]

patchwork/app.py

@@ -58,7 +58,10 @@ def list_option_callback(ctx: click.Context, param: click.Parameter, value: str
     ctx.exit()
 
 
-def find_patchflow(possible_module_paths: Iterable[str], patchflow: str) -> Any | None:
+def find_patchflow(possible_module_paths: Iterable[str], patchflow: str, trusted_modules: Iterable[str]) -> Any | None:
+    # Prepare whitelist set
+    trusted_modules_set = set(trusted_modules)
+    
     for module_path in possible_module_paths:
         try:
             spec = importlib.util.spec_from_file_location("custom_module", module_path)
@@ -71,14 +74,18 @@ def find_patchflow(possible_module_paths: Iterable[str], patchflow: str) -> Any
         except Exception:
             logger.debug(f"Patchflow {patchflow} not found as a file/directory in {module_path}")
 
-        try:
-            module = importlib.import_module(module_path)
-            logger.info(f"Patchflow {patchflow} loaded from {module_path}")
-            return getattr(module, patchflow)
-        except ModuleNotFoundError:
-            logger.debug(f"Patchflow {patchflow} not found as a module in {module_path}")
-        except AttributeError:
-            logger.debug(f"Patchflow {patchflow} not found in {module_path}")
+        # Check if module_path is trusted before importing it
+        if module_path in trusted_modules_set:
+            try:
+                module = importlib.import_module(module_path)
+                logger.info(f"Patchflow {patchflow} loaded from {module_path}")
+                return getattr(module, patchflow)
+            except ModuleNotFoundError:
+                logger.debug(f"Patchflow {patchflow} not found as a module in {module_path}")
+            except AttributeError:
+                logger.debug(f"Patchflow {patchflow} not found in {module_path}")
+        else:
+            logger.debug(f"Module path {module_path} is not in the list of trusted modules")
 
     return None