[pentest] add prompt injection pentesting

Author: baur-krykpayev

.gitignore

@@ -1,2 +1,3 @@
 .venv
 dist/
+__pycache__

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zenguard"
-version = "0.1.2"
+version = "0.1.3"
 description = "Plug-and-play production grade security for GenAI applications"
 authors = ["ZenGuard Team <hello@zenguard.ai>"]
 license = "MIT"
@@ -9,8 +9,14 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = "^3.9"
 httpx = "^0.27.0"
+tqdm = "^4.66.2"
 
 
+[tool.poetry.group.pentest.dependencies]
+openai = "^1.14.2"
+rapidfuzz = "^3.7.0"
+pandas = "^2.2.1"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

pyproject.toml

@@ -0,0 +1,11 @@
+import os
+
+from zenguard import Credentials, Detector, Endpoint, ZenGuard, ZenGuardConfig
+
+if __name__ == "__main__":
+    api_key = os.environ.get("ZEN_API_KEY")
+    if not api_key:
+        raise ValueError("ZEN_API_KEY is not set")
+    config = ZenGuardConfig(credentials=Credentials(api_key=api_key))
+    zenguard = ZenGuard(config=config)
+    zenguard.pentest(endpoint=Endpoint.ZENGUARD, detector=Detector.PROMPT_INJECTION)

tests/pentest.py

@@ -1 +1 @@
-from zenguard.zenguard import ZenGuard, ZenGuardConfig, Credentials, Detector
+from zenguard.zenguard import ZenGuard, ZenGuardConfig, Credentials, Detector, Endpoint

zenguard/__init__.py

@@ -0,0 +1 @@
+from zenguard.pentest.prompt_injections import config, prompt_data, prompting, run, scoring, visualization

zenguard/pentest/__init__.py

@@ -0,0 +1,26 @@
+import hashlib
+import json
+
+
+def hash_str(string):
+    return hashlib.md5(string.encode()).hexdigest()
+
+
+def hash_dict(d):
+    return hash_str(json.dumps(d))
+
+
+class DeepDict(dict):
+    def __missing__(self, key):
+        value = self[key] = type(self)()
+        return value
+
+    def __getitem__(self, key):
+        values = dict.__getitem__(self, key)
+        if isinstance(values, dict):
+            values = DeepDict(values)
+        if isinstance(values, list):
+            for i, v in enumerate(values):
+                if isinstance(v, dict):
+                    values[i] = DeepDict(v)
+        return values

zenguard/pentest/prompt_injections/__init__.py

@@ -0,0 +1,36 @@
+from . import prompt_data
+
+default_attack_config = {
+    "base": {
+        "text": prompt_data.ALL_OPENAI_PLAYGROUND_PROMPTS,
+    },
+    "attack": {
+        "text": (
+            prompt_data.goal_hikacking_attacks["ignore-print"],
+            prompt_data.goal_hikacking_attacks["ignore-say"],
+        ),
+        "rogue_string": prompt_data.rogue_strings["hate-humans"],
+        "scoring": "match_rogue_to_output",
+        "settings": {
+            "escape": prompt_data.escape_chars["n"],
+            "delimiter": prompt_data.delimiter_chars["dash"],
+            "escape_length": 10,
+            "escape_times": 2,
+        },
+    },
+    "config": {
+        "id": (0, 1, 2, 3),
+    },
+    "visualization": {
+        "columns": (
+            "config_model",
+            "config_temperature",
+            "prompt_instruction",
+            "attack_instruction",
+            "attack_rogue_string",
+            "config_id",
+            "score",
+        )
+    },
+}
+