diff --git a/patchwork/app.py b/patchwork/app.py
index 4149959a6..abd1a01bd 100644
--- a/patchwork/app.py
+++ b/patchwork/app.py
@@ -59,6 +59,9 @@ def list_option_callback(ctx: click.Context, param: click.Parameter, value: str
def find_patchflow(possible_module_paths: Iterable[str], patchflow: str) -> Any | None:
+ # Define a whitelist of allowed module paths
+ allowed_modules = {'allowed_module_1', 'allowed_module_2'}
+
for module_path in possible_module_paths:
try:
spec = importlib.util.spec_from_file_location("custom_module", module_path)
@@ -71,14 +74,18 @@ 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}")
+ # Check if the module is in the whitelist before importing
+ 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.warning(f"Module path {module_path} is not in the whitelist.")
return None
diff --git a/patchwork/common/tools/bash_tool.py b/patchwork/common/tools/bash_tool.py
index 8440f179a..1e802e781 100644
--- a/patchwork/common/tools/bash_tool.py
+++ b/patchwork/common/tools/bash_tool.py
@@ -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
+ 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:
return "Error: Command timed out after 60 seconds"
except Exception as e:
return f"Error: {str(e)}"
+
diff --git a/patchwork/common/tools/csvkit_tool.py b/patchwork/common/tools/csvkit_tool.py
index a1ef8dc59..9c88a88d9 100644
--- a/patchwork/common/tools/csvkit_tool.py
+++ b/patchwork/common/tools/csvkit_tool.py
@@ -117,11 +117,12 @@ def execute(self, files: list[str], query: str) -> str:
files_to_insert = []
if db_path.is_file():
with sqlite3.connect(str(db_path)) as conn:
+ cursor = conn.cursor()
for file in files:
- res = conn.execute(
- f"SELECT 1 from {file.removesuffix('.csv')}",
- )
- if res.fetchone() is None:
+ table_name = file.removesuffix('.csv')
+ cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
+ res = cursor.fetchone()
+ if res is None:
files_to_insert.append(file)
else:
files_to_insert = files
diff --git a/patchwork/common/tools/git_tool.py b/patchwork/common/tools/git_tool.py
new file mode 100644
index 000000000..d5acbfa05
--- /dev/null
+++ b/patchwork/common/tools/git_tool.py
@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import os
+import subprocess
+
+from patchwork.common.tools.tool import Tool
+
+
+class GitTool(Tool, tool_name="git_tool", abc_register=False):
+ def __init__(self, path: str):
+ super().__init__()
+ self.path = path
+
+ @property
+ def json_schema(self) -> dict:
+ return {
+ "name": "git_tool",
+ "description": """\
+Access to the Git CLI, the command is also `git` all args provided are used as is.
+""",
+ "input_schema": {
+ "type": "object",
+ "properties": {
+ "args": {
+ "type": "array",
+ "items": {"type": "string"},
+ "description": """
+The args to run `git` command with.
+E.g.
+[\"commit\", \"-m\", \"A commit message\"] to commit changes with a commit message.
+[\"add\", \".\"] to stage all changed files.
+""",
+ }
+ },
+ "required": ["args"],
+ },
+ }
+
+ def execute(self, args: list[str]) -> str:
+ env = os.environ.copy()
+ p = subprocess.run(
+ ["git", *args],
+ env=env,
+ cwd=self.path,
+ text=True,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ )
+ return p.stdout
diff --git a/patchwork/common/tools/github_tool.py b/patchwork/common/tools/github_tool.py
index aa5d5effe..16e418c08 100644
--- a/patchwork/common/tools/github_tool.py
+++ b/patchwork/common/tools/github_tool.py
@@ -6,7 +6,7 @@
from patchwork.common.tools.tool import Tool
-class GitHubTool(Tool, tool_name="github_tool"):
+class GitHubTool(Tool, tool_name="github_tool", abc_register=False):
def __init__(self, path: str, gh_token: str):
super().__init__()
self.path = path
diff --git a/patchwork/common/utils/dependency.py b/patchwork/common/utils/dependency.py
index 27b89bfed..6e62f0f92 100644
--- a/patchwork/common/utils/dependency.py
+++ b/patchwork/common/utils/dependency.py
@@ -6,9 +6,13 @@
"notification": ["slack_sdk"],
}
+__ALLOWED_MODULES = {module for modules in __DEPENDENCY_GROUPS.values() for module in modules}
@lru_cache(maxsize=None)
def import_with_dependency_group(name):
+ if name not in __ALLOWED_MODULES:
+ raise ImportError(f"Import of untrusted module '{name}' is not allowed.")
+
try:
return importlib.import_module(name)
except ImportError:
@@ -20,6 +24,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..0c94790e7 100644
--- a/patchwork/common/utils/step_typing.py
+++ b/patchwork/common/utils/step_typing.py
@@ -106,9 +106,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]]:
+ allowed_modules = {"module_1.typed", "module_2.typed"} # Example whitelist
module_path, _, _ = step.__module__.rpartition(".")
step_name = step.__name__
- type_module = importlib.import_module(f"{module_path}.typed")
+ type_module_path = f"{module_path}.typed"
+
+ if type_module_path not in allowed_modules:
+ raise ValueError(f"Importing from {type_module_path} is not allowed")
+
+ type_module = importlib.import_module(type_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:
diff --git a/patchwork/steps/CallShell/CallShell.py b/patchwork/steps/CallShell/CallShell.py
index 98ee55a74..720328b33 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)
+ command_args = shlex.split(self.script)
+ p = subprocess.run(command_args, 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 bc8d319c1..ede41c109 100644
--- a/patchwork/steps/GitHubAgent/GitHubAgent.py
+++ b/patchwork/steps/GitHubAgent/GitHubAgent.py
@@ -5,6 +5,7 @@
AgentConfig,
AgenticStrategyV2,
)
+from patchwork.common.tools.git_tool import GitTool
from patchwork.common.tools.github_tool import GitHubTool
from patchwork.common.utils.utils import mustache_render
from patchwork.step import Step
@@ -34,10 +35,13 @@ def __init__(self, inputs):
AgentConfig(
name="Assistant",
model="gemini-2.0-flash",
- tool_set=dict(github_tool=GitHubTool(base_path, inputs["github_api_key"])),
+ tool_set=dict(
+ github_tool=GitHubTool(base_path, inputs["github_api_key"]),
+ git_tool=GitTool(base_path),
+ ),
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 can access github through the `gh` CLI app through the `github_tool`, and `git` through the `git_tool`.
Your `gh` app has already been authenticated.
""",
)
diff --git a/pyproject.toml b/pyproject.toml
index ab89dee54..d136584c5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[tool.poetry]
name = "patchwork-cli"
-version = "0.0.123"
+version = "0.0.124"
description = ""
authors = ["patched.codes"]
license = "AGPL"