PatchWork AutoFix

patchwork/app.py

@@ -216,14 +216,16 @@ 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}")
+        valid_modules = ['module1', 'module2']  # Add allowed modules to the whitelist
+        if module_path in valid_modules:
+            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}")
 
     return None
 

patchwork/common/utils/dependency.py

@@ -7,9 +7,13 @@
     "notification": ["slack_sdk"],
 }
 
+ALLOWED_MODULES = set(__DEPENDENCY_GROUPS.values())
 
 @lru_cache(maxsize=None)
 def import_with_dependency_group(name):
+    if name not in ALLOWED_MODULES:
+        raise ImportError(f"Module {name} is not whitelisted")
+
     try:
         return importlib.import_module(name)
     except ImportError:

patchwork/common/utils/step_typing.py

@@ -119,7 +119,7 @@ def validate_step_with_inputs(input_keys: Set[str], step: Type[Step]) -> Tuple[S
     step_report = {}
     for key in step_input_model.__required_keys__:
         if key not in input_keys:
-            step_report[key] = f"Missing required input data"
+            step_report[key] = "Missing required input data"
             continue
 
     step_type_hints = get_type_hints(step_input_model, include_extras=True)
@@ -129,7 +129,7 @@ def validate_step_with_inputs(input_keys: Set[str], step: Type[Step]) -> Tuple[S
             continue
 
         if key in step_report.keys():
-            step_report[key] = step_type_config.msg or f"Missing required input data"
+            step_report[key] = step_type_config.msg or "Missing required input data"
             continue
 
         is_ok, msg = validate_step_type_config_with_inputs(key, input_keys, step_type_config)