PatchWork AutoFix

patchwork/app.py

@@ -59,7 +59,13 @@ def list_option_callback(ctx: click.Context, param: click.Parameter, value: str
 
 
 def find_patchflow(possible_module_paths: Iterable[str], patchflow: str) -> Any | None:
+    whitelisted_modules = {"allowed_module1", "allowed_module2"}  # Adjust whitelist as necessary
+
     for module_path in possible_module_paths:
+        if module_path not in whitelisted_modules:
+            logger.debug(f"Module path {module_path} is not whitelisted.")
+            continue
+
         try:
             spec = importlib.util.spec_from_file_location("custom_module", module_path)
             module = importlib.util.module_from_spec(spec)

patchwork/common/tools/bash_tool.py

@@ -2,12 +2,12 @@
 
 import subprocess
 from pathlib import Path
+import shlex
 
 from typing_extensions import Optional
 
 from patchwork.common.tools.tool import Tool
 
-
 class BashTool(Tool, tool_name="bash"):
     def __init__(self, path: Path):
         super().__init__()
@@ -45,10 +45,11 @@ def execute(
 
         try:
             result = subprocess.run(
-                command, shell=True, cwd=self.path, capture_output=True, text=True, timeout=60  # Add timeout for safety
+                shlex.split(command), shell=False, cwd=self.path, capture_output=True, text=True, timeout=60  # Add timeout for safety
             )
             return result.stdout if result.returncode == 0 else f"Error: {result.stderr}"
         except subprocess.TimeoutExpired:
             return "Error: Command timed out after 60 seconds"
         except Exception as e:
             return f"Error: {str(e)}"
+

patchwork/common/tools/csvkit_tool.py

@@ -118,8 +118,9 @@ 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')
                     res = conn.execute(
-                        f"SELECT 1 from {file.removesuffix('.csv')}",
+                        "SELECT 1 FROM ?", (table_name,)
                     )
                     if res.fetchone() is None:
                         files_to_insert.append(file)

patchwork/common/utils/dependency.py

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

patchwork/common/utils/step_typing.py

@@ -108,7 +108,15 @@ 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]]:
     module_path, _, _ = step.__module__.rpartition(".")
     step_name = step.__name__
-    type_module = importlib.import_module(f"{module_path}.typed")
+
+    # Whitelist of allowed modules
+    allowed_modules = {"allowed_module1.typed", "allowed_module2.typed"}
+    full_module_path = f"{module_path}.typed"
+
+    if full_module_path not in allowed_modules:
+        raise ImportError(f"Attempt to import disallowed module {full_module_path}")
+
+    type_module = importlib.import_module(full_module_path)
     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:

patchwork/steps/AgenticLLMV2/AgenticLLMV2.py

@@ -18,15 +18,15 @@ def __init__(self, inputs):
             base_path = str(Path.cwd())
         self.conversation_limit = int(inputs.get("max_agent_calls", 1))
         self.agentic_strategy = AgenticStrategyV2(
-            model="claude-3-5-sonnet-latest",
+            model=inputs.get("strategy_model", "claude-3-5-sonnet-latest"),
             llm_client=AioLlmClient.create_aio_client(inputs),
             template_data=inputs.get("prompt_value", {}),
             system_prompt_template=inputs.get("system_prompt", "Summarise from our previous conversation"),
             user_prompt_template=inputs.get("user_prompt"),
             agent_configs=[
                 AgentConfig(
                     name="Assistant",
-                    model="claude-3-7-sonnet-latest",
+                    model=inputs.get("agent_model", "claude-3-7-sonnet-latest"),
                     tool_set=Tool.get_tools(path=base_path),
                     system_prompt=inputs.get("agent_system_prompt"),
                 )

patchwork/steps/AgenticLLMV2/typed.py

@@ -9,9 +9,32 @@ class AgenticLLMV2Inputs(TypedDict, total=False):
     system_prompt: str
     user_prompt: str
     max_agent_calls: Annotated[int, StepTypeConfig(is_config=True)]
-    anthropic_api_key: str
+    strategy_model: str
+    agent_model: str
     agent_system_prompt: str
     example_json: str
+    openai_api_key: Annotated[
+        str,
+        StepTypeConfig(
+            is_config=True, or_op=["patched_api_key", "google_api_key", "client_is_gcp", "anthropic_api_key"]
+        ),
+    ]
+    anthropic_api_key: Annotated[
+        str,
+        StepTypeConfig(is_config=True, or_op=["patched_api_key", "google_api_key", "client_is_gcp", "openai_api_key"]),
+    ]
+    google_api_key: Annotated[
+        str,
+        StepTypeConfig(
+            is_config=True, or_op=["patched_api_key", "openai_api_key", "client_is_gcp", "anthropic_api_key"]
+        ),
+    ]
+    client_is_gcp: Annotated[
+        str,
+        StepTypeConfig(
+            is_config=True, or_op=["patched_api_key", "openai_api_key", "anthropic_api_key", "google_api_key"]
+        ),
+    ]
 
 
 class AgenticLLMV2Outputs(TypedDict):

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)
+        args = shlex.split(self.script)
+        p = subprocess.run(args, shell=False, capture_output=True, text=True, cwd=self.working_dir, env=self.env)
         try:
             p.check_returncode()
         except subprocess.CalledProcessError as e:
@@ -57,3 +58,4 @@ def run(self) -> dict:
         logger.info(f"stdout: \n{p.stdout}")
         logger.info(f"stderr:\n{p.stderr}")
         return dict(stdout_output=p.stdout, stderr_output=p.stderr)
+

patchwork/steps/GitHubAgent/GitHubAgent.py

@@ -21,7 +21,7 @@ def __init__(self, inputs):
             model="claude-3-5-sonnet-latest",
             llm_client=AioLlmClient.create_aio_client(inputs),
             template_data=dict(),
-            system_prompt_template=f"""\
+            system_prompt_template="""\
 Please summarise the conversation given and provide the result in the structure that is asked of you.
 """,
             user_prompt_template=f"""\
@@ -34,15 +34,17 @@ def __init__(self, inputs):
                 AgentConfig(
                     name="Assistant",
                     model="gemini-2.0-flash",
-                    tool_set=dict(github_tool=GitHubTool(base_path, inputs["github_api_token"])),
+                    tool_set=dict(github_tool=GitHubTool(base_path, inputs["github_api_key"])),
                     system_prompt="""\
-You are a senior software developer helping the program manager to obtain some data from GitHub. 
-You can access github through the `gh` CLI app. 
+You are a senior software developer helping the program manager to obtain some data from GitHub.
+You can access github through the `gh` CLI app.
 Your `gh` app has already been authenticated.
 """,
                 )
             ],
-            example_json=inputs.get("example_json"),
+            example_json=inputs.get(
+                "example_json", '{"summary_of_actions": "1. Retrieved the list of repositories. 2. ..."}'
+            ),
         )
 
     def run(self) -> dict:

patchwork/steps/GitHubAgent/typed.py

@@ -1,16 +1,18 @@
-from typing_extensions import Annotated, Any, Dict, TypedDict
+from typing_extensions import Annotated, Any, Dict, Optional, TypedDict
 
 from patchwork.common.utils.step_typing import StepTypeConfig
 
 
 class __GitHubAgentRequiredInputs(TypedDict):
+    github_api_key: str
     task: str
 
 
 class GitHubAgentInputs(__GitHubAgentRequiredInputs, total=False):
     base_path: str
     prompt_value: Dict[str, Any]
     max_llm_calls: Annotated[int, StepTypeConfig(is_config=True)]
+    example_json: str
     openai_api_key: Annotated[
         str,
         StepTypeConfig(