diff --git a/patchwork/app.py b/patchwork/app.py index 4149959a6..18c30f842 100644 --- a/patchwork/app.py +++ b/patchwork/app.py @@ -59,6 +59,7 @@ def list_option_callback(ctx: click.Context, param: click.Parameter, value: str def find_patchflow(possible_module_paths: Iterable[str], patchflow: str) -> Any | None: + allowed_modules = {"module1", "module2", "module3"} # whitelist of allowed modules for module_path in possible_module_paths: try: spec = importlib.util.spec_from_file_location("custom_module", module_path) @@ -71,14 +72,17 @@ 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}") + if module_path in allowed_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}") + else: + logger.debug(f"Module {module_path} is not in the allowed list and will not be imported.") return None diff --git a/patchwork/common/tools/bash_tool.py b/patchwork/common/tools/bash_tool.py index 8440f179a..78126b97b 100644 --- a/patchwork/common/tools/bash_tool.py +++ b/patchwork/common/tools/bash_tool.py @@ -45,7 +45,7 @@ def execute( try: result = subprocess.run( - command, shell=True, cwd=self.path, capture_output=True, text=True, timeout=60 # Add timeout for safety + command.split(), shell=False, cwd=self.path, capture_output=True, text=True, timeout=60 ) return result.stdout if result.returncode == 0 else f"Error: {result.stderr}" except subprocess.TimeoutExpired: diff --git a/patchwork/common/tools/csvkit_tool.py b/patchwork/common/tools/csvkit_tool.py index a1ef8dc59..7d4278cfb 100644 --- a/patchwork/common/tools/csvkit_tool.py +++ b/patchwork/common/tools/csvkit_tool.py @@ -118,8 +118,11 @@ def execute(self, files: list[str], query: str) -> str: if db_path.is_file(): with sqlite3.connect(str(db_path)) as conn: for file in files: + table_name = file.removesuffix('.csv') + if not re.match(r'^[A-Za-z0-9_]+$', table_name): + continue res = conn.execute( - f"SELECT 1 from {file.removesuffix('.csv')}", + "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (table_name,) ) if res.fetchone() is None: files_to_insert.append(file) diff --git a/patchwork/common/utils/dependency.py b/patchwork/common/utils/dependency.py index 27b89bfed..2eeea3764 100644 --- a/patchwork/common/utils/dependency.py +++ b/patchwork/common/utils/dependency.py @@ -6,9 +6,11 @@ "notification": ["slack_sdk"], } - @lru_cache(maxsize=None) def import_with_dependency_group(name): + allowed_modules = {mod for mods in __DEPENDENCY_GROUPS.values() for mod in mods} + if name not in allowed_modules: + raise ImportError(f"Module {name} is not allowed to be imported.") try: return importlib.import_module(name) except ImportError: @@ -20,6 +22,5 @@ def import_with_dependency_group(name): error_msg = f"Please `pip install patchwork-cli[{dependency_group}]` to use this step" raise ImportError(error_msg) - def slack_sdk(): return import_with_dependency_group("slack_sdk") diff --git a/patchwork/common/utils/step_typing.py b/patchwork/common/utils/step_typing.py index d349f7fc1..f21dbe54e 100644 --- a/patchwork/common/utils/step_typing.py +++ b/patchwork/common/utils/step_typing.py @@ -106,9 +106,13 @@ def validate_step_type_config_with_inputs( def validate_step_with_inputs(input_keys: Set[str], step: Type[Step]) -> Tuple[Set[str], Dict[str, str]]: + allowed_modules = ["module1.typed", "module2.typed"] # Example whitelist module_path, _, _ = step.__module__.rpartition(".") step_name = step.__name__ - type_module = importlib.import_module(f"{module_path}.typed") + type_module_name = f"{module_path}.typed" + if type_module_name not in allowed_modules: + raise ImportError(f"Module {type_module_name} is not allowed to be imported") + type_module = importlib.import_module(type_module_name) step_input_model = getattr(type_module, f"{step_name}Inputs", __NOT_GIVEN) step_output_model = getattr(type_module, f"{step_name}Outputs", __NOT_GIVEN) if step_input_model is __NOT_GIVEN: diff --git a/patchwork/steps/CallShell/CallShell.py b/patchwork/steps/CallShell/CallShell.py index 98ee55a74..56dd511b0 100644 --- a/patchwork/steps/CallShell/CallShell.py +++ b/patchwork/steps/CallShell/CallShell.py @@ -46,7 +46,8 @@ def __parse_env_text(env_text: str) -> dict[str, str]: return env def run(self) -> dict: - p = subprocess.run(self.script, shell=True, capture_output=True, text=True, cwd=self.working_dir, env=self.env) + script_list = shlex.split(self.script) + p = subprocess.run(script_list, shell=False, capture_output=True, text=True, cwd=self.working_dir, env=self.env) try: p.check_returncode() except subprocess.CalledProcessError as e: diff --git a/patchwork/steps/GitHubAgent/GitHubAgent.py b/patchwork/steps/GitHubAgent/GitHubAgent.py index e431cfd70..bc8d319c1 100644 --- a/patchwork/steps/GitHubAgent/GitHubAgent.py +++ b/patchwork/steps/GitHubAgent/GitHubAgent.py @@ -18,7 +18,7 @@ def __init__(self, inputs): data = inputs.get("prompt_value", {}) task = mustache_render(inputs["task"], data) self.agentic_strategy = AgenticStrategyV2( - model="claude-3-5-sonnet-latest", + model="gemini-2.0-flash", llm_client=AioLlmClient.create_aio_client(inputs), template_data=dict(), system_prompt_template="""\ diff --git a/pyproject.toml b/pyproject.toml index 17a4adc1c..ab89dee54 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "patchwork-cli" -version = "0.0.122" +version = "0.0.123" description = "" authors = ["patched.codes"] license = "AGPL"