PatchWork AutoFix

patchwork/app.py

@@ -59,6 +59,8 @@ def list_option_callback(ctx: click.Context, param: click.Parameter, value: str
 
 
 def find_patchflow(possible_module_paths: Iterable[str], patchflow: str) -> Any | None:
+    trusted_modules = {"trusted_module1", "trusted_module2", "trusted_module3"}
+
     for module_path in possible_module_paths:
         try:
             spec = importlib.util.spec_from_file_location("custom_module", module_path)
@@ -71,14 +73,15 @@ 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 trusted_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}")
 
     return None
 

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)}"
+

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 ? LIMIT 1", (table_name,)
                     )
                     if res.fetchone() is None:
                         files_to_insert.append(file)

patchwork/common/utils/dependency.py

@@ -6,9 +6,12 @@
     "notification": ["slack_sdk"],
 }
 
+__ALLOWED_MODULES = set(sum(__DEPENDENCY_GROUPS.values(), []))
 
 @lru_cache(maxsize=None)
 def import_with_dependency_group(name):
+    if name not in __ALLOWED_MODULES:
+        raise ImportError(f"Importing {name} is not allowed.")
     try:
         return importlib.import_module(name)
     except ImportError:

patchwork/common/utils/step_typing.py

@@ -108,6 +108,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]]:
     module_path, _, _ = step.__module__.rpartition(".")
     step_name = step.__name__
+    # Define a whitelist of allowed module paths for importing
+    allowed_module_paths = {'approved.module1', 'approved.module2'} 
+
+    # Add a check against the whitelist
+    if module_path not in allowed_module_paths:
+        raise ValueError(f"Module path '{module_path}' is not allowed")
+
     type_module = importlib.import_module(f"{module_path}.typed")
     step_input_model = getattr(type_module, f"{step_name}Inputs", __NOT_GIVEN)
     step_output_model = getattr(type_module, f"{step_name}Outputs", __NOT_GIVEN)

patchwork/steps/CallShell/CallShell.py

@@ -46,7 +46,7 @@ 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)
+        p = subprocess.run(shlex.split(self.script), shell=False, capture_output=True, text=True, cwd=self.working_dir, env=self.env)
         try:
             p.check_returncode()
         except subprocess.CalledProcessError as e:

pyproject.toml

@@ -90,7 +90,7 @@ all = [
 
 [tool.poetry.group.dev.dependencies]
 setuptools = "*"
-poethepoet = { version = "^0.27.0", extras = ["poetry-plugin"] }
+poethepoet = {version = "^0.33.1", extras = ["poetry-plugin"]}
 mypy = "^1.7.1"
 types-requests = "~2.31.0"
 black = "^23.12.0"