diff --git a/patchwork/app.py b/patchwork/app.py index ae6df6e5f..8510090ad 100644 --- a/patchwork/app.py +++ b/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 diff --git a/patchwork/common/utils/dependency.py b/patchwork/common/utils/dependency.py index 040822136..7ab1a9475 100644 --- a/patchwork/common/utils/dependency.py +++ b/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: diff --git a/patchwork/common/utils/step_typing.py b/patchwork/common/utils/step_typing.py index d349f7fc1..2477cc759 100644 --- a/patchwork/common/utils/step_typing.py +++ b/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: