PatchWork AutoFix

patchwork/app.py

@@ -216,15 +216,6 @@ 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}")
-
     return None
 
 

patchwork/common/utils/dependency.py

@@ -7,9 +7,14 @@
     "notification": ["slack_sdk"],
 }
 
+def is_valid_dependency(name):
+    return name in __DEPENDENCY_GROUPS.values()
 
 @lru_cache(maxsize=None)
 def import_with_dependency_group(name):
+    if not is_valid_dependency(name):
+        raise ImportError(f"Invalid dependency {name}")
+
     try:
         return importlib.import_module(name)
     except ImportError:

patchwork/common/utils/step_typing.py

@@ -105,8 +105,13 @@ def validate_step_type_config_with_inputs(
     return True, step_type_config.msg
 
 
+WHITELISTED_MODULES = {'allowed_module_1', 'allowed_module_2'} # Add allowed modules to the whitelist
+
 def validate_step_with_inputs(input_keys: Set[str], step: Type[Step]) -> Tuple[Set[str], Dict[str, str]]:
     module_path, _, _ = step.__module__.rpartition(".")
+    if module_path not in WHITELISTED_MODULES:
+        raise ValueError(f"Module {module_path} is not in the whitelist")
+
     step_name = step.__name__
     type_module = importlib.import_module(f"{module_path}.typed")
     step_input_model = getattr(type_module, f"{step_name}Inputs", __NOT_GIVEN)
@@ -119,7 +124,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 +134,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)
@@ -139,6 +144,7 @@ def validate_step_with_inputs(input_keys: Set[str], step: Type[Step]) -> Tuple[S
     return set(step_output_model.__required_keys__), step_report
 
 
+
 def find_step_type_config(python_type: type) -> Optional[StepTypeConfig]:
     annotated = find_annotated(python_type)
     if annotated is None: