add dedicated step

Author: CTY-git

patchwork/step.py

@@ -45,10 +45,7 @@ def __init__(self, inputs: DataPoint):
         """
 
         # check if the inputs have the required keys
-        if self.__input_class is not None:
-            missing_keys = self.__input_class.__required_keys__.difference(inputs.keys())
-            if len(missing_keys) > 0:
-                raise ValueError(f"Missing required data: {list(missing_keys)}")
+        self.test_inputs(inputs)
 
         # store the inputs
         self.inputs = inputs
@@ -77,6 +74,12 @@ def __init_subclass__(cls, **kwargs):
             cls.__output_class = output_class
         else:
             cls.__output_class = None
+    @classmethod
+    def test_inputs(cls, inputs: DataPoint):
+        if cls.__input_class is not None:
+            missing_keys = cls.__input_class.__required_keys__.difference(inputs.keys())
+            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)

patchwork/steps/CallCommand/CallCommand.py

@@ -1,4 +1,5 @@
-import json
+from __future__ import annotations
+
 import shlex
 import subprocess
 from pathlib import Path
@@ -17,12 +18,12 @@ def __init__(self, inputs: dict):
         self.env = self.__parse_env_text(inputs.get("env", ""))
 
     @staticmethod
-    def __parse_env_text(env_text):
+    def __parse_env_text(env_text: str) -> dict[str, str]:
         env_spliter = shlex.shlex(env_text, posix=True)
         env_spliter.whitespace_split = True
         env_spliter.whitespace += ";"
 
-        env = dict()
+        env: dict[str, str] = dict()
         for env_assign in env_spliter:
             env_assign_spliter = shlex.shlex(env_assign, posix=True)
             env_assign_spliter.whitespace_split = True

patchwork/steps/ScanPSFuzz/ScanPSFuzz.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import subprocess
+
+from patchwork.logger import logger
+from patchwork.steps import CallCommand
+from patchwork.steps.ScanPSFuzz.typed import ScanPSFuzzInputs, ScanPSFuzzOutputs
+
+
+class ScanPSFuzz(CallCommand, input_class=ScanPSFuzzInputs, output_class=ScanPSFuzzOutputs):
+    def __init__(self, inputs: dict):
+        if not self.__is_ps_fuzz_installed():
+            raise ValueError("""\
+`prompt-security-fuzzer` is not installed. Please install with the following instructions:
+1. Install pipx: https://github.com/pypa/pipx
+2. pipx install prompt-security-fuzzer
+3. pipx inject prompt-security-fuzzer setuptools
+""")
+        self.test_inputs(inputs)
+        wrapped_input = dict(
+            command="prompt-security-fuzzer",
+            command_args=f'-b {inputs["prompt_file_path"]}',
+            env=f'OPENAI_API_KEY={inputs["openai_api_key"]}'
+        )
+
+        working_dir = inputs.get("working_dir"),
+        if working_dir is not None:
+            wrapped_input["working_dir"] = working_dir
+
+        super().__init__(wrapped_input)
+
+    def __is_ps_fuzz_installed(self):
+        try:
+            subprocess.run(["prompt-security-fuzzer", "-h"], check=True)
+            return True
+        except subprocess.CalledProcessError as e:
+            err = e
+        except FileNotFoundError as e:
+            err = e
+        # If the command fails, prompt-security-fuzzer is not installed
+        logger.info(f"prompt-security-fuzzer is not installed: {err}")
+        return False
+

patchwork/steps/ScanPSFuzz/__init__.py

@@ -0,0 +1,15 @@
+from typing_extensions import Annotated, TypedDict
+
+from patchwork.common.utils.step_typing import StepTypeConfig
+
+
+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)]
+
+
+class ScanPSFuzzOutputs(TypedDict):
+    stdout_output: str

patchwork/steps/ScanPSFuzz/typed.py

@@ -42,6 +42,7 @@
 from patchwork.steps.ReadPRDiffs.ReadPRDiffs import ReadPRDiffs
 from patchwork.steps.ReadPRs.ReadPRs import ReadPRs
 from patchwork.steps.ScanDepscan.ScanDepscan import ScanDepscan
+from patchwork.steps.ScanPSFuzz.ScanPSFuzz import ScanPSFuzz
 from patchwork.steps.ScanSemgrep.ScanSemgrep import ScanSemgrep
 from patchwork.steps.SimplifiedLLM.SimplifiedLLM import SimplifiedLLM
 from patchwork.steps.SimplifiedLLMOnce.SimplifiedLLMOnce import SimplifiedLLMOnce
@@ -90,6 +91,7 @@
     "ReadPRDiffsPB",
     "ReadPRs",
     "ScanDepscan",
+    "ScanPSFuzz"
     "ScanSemgrep",
     "SimplifiedLLM",
     "SimplifiedLLMOnce",