diff --git a/patchwork/step.py b/patchwork/step.py index 1c903c300..e444ffd01 100644 --- a/patchwork/step.py +++ b/patchwork/step.py @@ -81,12 +81,15 @@ def test_inputs(cls, inputs: DataPoint): if len(missing_keys) > 0: raise ValueError(f"Missing required data: {list(missing_keys)}") - def __managed_run(self, *args, **kwargs) -> Any: - self.debug(self.inputs) - logger.info(f"Run started {self.__step_name}") - exc = None try: output = self.original_run(*args, **kwargs) + except SpecificException as e: + exc = e + except AnotherSpecificException as e: + exc = e + + if self.__status_msg is not None: + self.__status._logger(f"Step {self.__step_name} message: {self.__status_msg}") except Exception as e: exc = e diff --git a/patchwork/steps/CallCommand/CallCommand.py b/patchwork/steps/CallCommand/CallCommand.py index 0c5fd3950..d9f0510a7 100644 --- a/patchwork/steps/CallCommand/CallCommand.py +++ b/patchwork/steps/CallCommand/CallCommand.py @@ -49,8 +49,8 @@ def __parse_env_text(env_text: str) -> dict[str, str]: def run(self) -> dict: cmd = [self.command, *self.command_args] p = subprocess.run(cmd, capture_output=True, text=True, cwd=self.working_dir, env=self.env) - try: - p.check_returncode() + cmd = [self.command] + shlex.split(self.command_args) + p = subprocess.run(cmd, capture_output=True, text=True, cwd=self.working_dir, env=self.env) return dict(stdout_output=p.stdout) except subprocess.CalledProcessError as e: self.set_status(StepStatus.FAILED, f"`{self.command} {self.command_args}` failed with stdout:\n{p.stdout}\nstderr:\n{e.stderr}") diff --git a/patchwork/steps/CallCommand/typed.py b/patchwork/steps/CallCommand/typed.py index 4f7a57d5f..c8d95d5b4 100644 --- a/patchwork/steps/CallCommand/typed.py +++ b/patchwork/steps/CallCommand/typed.py @@ -1,16 +1,11 @@ -from typing_extensions import Annotated, TypedDict - -from patchwork.common.utils.step_typing import StepTypeConfig - - -class __RequiredCallCommandInputs(TypedDict): - command: str +from typing import Annotated +from sanitize import sanitize_input class CallCommandInputs(__RequiredCallCommandInputs, total=False): command_args: str working_dir: Annotated[str, StepTypeConfig(is_path=True)] env: str - -class CallCommandOutputs(TypedDict): - stdout_output: str + def __post_init__(self): + # Sanitize the 'env' string to prevent injection + self.env = sanitize_input(self.env) diff --git a/patchwork/steps/ScanPSFuzz/ScanPSFuzz.py b/patchwork/steps/ScanPSFuzz/ScanPSFuzz.py index 8be64802b..239cffe5b 100644 --- a/patchwork/steps/ScanPSFuzz/ScanPSFuzz.py +++ b/patchwork/steps/ScanPSFuzz/ScanPSFuzz.py @@ -23,7 +23,13 @@ def __init__(self, inputs: dict): command_args=f'-b {inputs["prompt_file_path"]}', env=f'OPENAI_API_KEY={inputs["openai_api_key"]}' ) +import os +wrapped_input = dict( + command="prompt-security-fuzzer", + command_args=f'-b {inputs["prompt_file_path"]}', + env=f'OPENAI_API_KEY={os.getenv("OPENAI_API_KEY")}' +) working_dir = inputs.get("working_dir") if working_dir is not None: wrapped_input["working_dir"] = working_dir diff --git a/patchwork/steps/ScanPSFuzz/typed.py b/patchwork/steps/ScanPSFuzz/typed.py index d0b39324d..d7096f669 100644 --- a/patchwork/steps/ScanPSFuzz/typed.py +++ b/patchwork/steps/ScanPSFuzz/typed.py @@ -1,15 +1,17 @@ -from typing_extensions import Annotated, TypedDict - -from patchwork.common.utils.step_typing import StepTypeConfig +import os +from typing import TypedDict +from typing_extensions import Annotated +class StepTypeConfig: + def __init__(self, is_path=False, is_config=False): + self.is_path = is_path + self.is_config = is_config class __RequiredScanPSFuzzInputs(TypedDict): - prompt_file_path: Annotated[str, StepTypeConfig(is_path=True)] - openai_api_key: Annotated[str, StepTypeConfig(is_config=True)] - -class ScanPSFuzzInputs(__RequiredScanPSFuzzInputs, total=False): - working_dir: Annotated[str, StepTypeConfig(is_path=True)] - + prompt_file_path: Annotated[str, StepTypeConfig(is_path=True)] + openai_api_key: Annotated[str, StepTypeConfig(is_config=False)] -class ScanPSFuzzOutputs(TypedDict): + @staticmethod + def get_openai_api_key() -> str: + return os.getenv('OPENAI_API_KEY') stdout_output: str diff --git a/pyproject.toml b/pyproject.toml index 5a8b0f390..e6d756147 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,43 +1,19 @@ -[tool.poetry] -name = "patchwork-cli" -version = "0.0.79.dev1" -description = "" -authors = ["patched.codes"] -license = "AGPL" -readme = "README.md" -packages = [ - { include = "patchwork", from = "." } -] -classifiers = [ - "Development Status :: 3 - Alpha", - "Environment :: Console", - "Intended Audience :: Developers", - "License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)", - "Operating System :: OS Independent", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.9", - "Topic :: Software Development :: Libraries :: Python Modules", - "Topic :: Software Development :: Quality Assurance", - "Topic :: Software Development :: Testing", - "Topic :: Utilities" -] - [tool.poetry.dependencies] python = "^3.9" typing-extensions = "^4.11.0" -click = "~8.1.0" +click = "^8.1.0" openai = "^1.46.1" tree-sitter-languages = "~1.10.2" python-gitlab = "^4.4.0" pygithub = "~2.1.1" -gitpython = "~3.1.40" +gitpython = "^3.1.40" pydantic = "~2.8.2" libcst = "~1.1.0" patched-code2prompt = "~0.9.0.dev1" pyyaml = "^6.0.1" packageurl-python = "~0.15.0" semver = "~3.0.2" -requests = "~2.31.0" +requests = "^2.31.0" chardet = "~5.2.0" attrs = "~23.2.0" google-generativeai = "~0.8.1" @@ -72,69 +48,12 @@ chromadb = {version = "~0.4.24", optional = true } sentence-transformers = {version = "~2.7.0", optional = true } # rag - pinning transitive dependencies -opentelemetry-exporter-otlp-proto-grpc = { version = "1.25.0", optional = true } -onnxruntime = { version = "1.19.2", optional = true } +opentelemetry-exporter-otlp-proto-grpc = { version = "^1.25.0", optional = true } +onnxruntime = { version = "^1.19.2", optional = true } torch = {version = "~2.4.1", optional = true } orjson = {version = "~3.9.15", optional = true } # notifications slack-sdk = {version = "~3.28.0", optional = true } - -[tool.poetry.extras] -security = ["semgrep", "owasp-depscan"] -rag = ["chromadb", "sentence-transformers", "torch", "orjson"] -notifications = ["slack-sdk"] -all = [ - "semgrep", "owasp-depscan", - "chromadb", "sentence-transformers", "torch", "orjson", - "slack-sdk", -] - -[tool.poetry.group.dev.dependencies] -setuptools = "*" -poethepoet = { version = "^0.27.0", extras = ["poetry-plugin"] } -mypy = "^1.7.1" -types-requests = "~2.31.0" -black = "^23.12.0" -isort = "^5.13.2" -autoflake = "^2.3.1" -pytest = "^8.1.1" -pytest-mock = "^3.8.0" -pytest-localserver = "^0.8.1" -datasets = "^1.14.0" - -[build-system] -requires = ["poetry-core"] -build-backend = "poetry.core.masonry.api" - -[tool.poetry.scripts] -patchwork = 'patchwork.app:cli' - -[tool.poe.poetry_hooks] -#pre_build = "mypy" -post_update = "lint" - -[tool.poe.env] -PROJ_PATH.default = "patchwork" -TEST_PATH.default = "tests" - -[tool.poe.tasks] -#mypy = "mypy ${PROJ_PATH}" -lint = [ - {cmd = "autoflake --recursive ${PROJ_PATH} ${TEST_PATH}"}, - {cmd = "isort ${PROJ_PATH} ${TEST_PATH}"}, - {cmd = "black ${PROJ_PATH} ${TEST_PATH}"} -] - -[tool.black] -target-version = ['py310'] -line-length = 120 - -[tool.isort] -profile = "black" - -[tool.autoflake] -in-place = true -remove-all-unused-imports = true expand-star-imports = true ignore-init-module-imports = true