[pentest] pentest OpenAI

Author: baur-krykpayev

README.md

@@ -37,6 +37,32 @@ Test the capabilities of ZenGuard AI in our ZenGuard [Playground](https://consol
 
 A more detailed documentation is available at [docs.zenguard.ai](https://docs.zenguard.ai/).
 
+
+# Pentesting
+
+You can run pentest against both ZenGuard AI and (optionally) ChatGPT.
+
+Clone this repo and install requirements.
+
+Run pentest against ZenGuard AI:
+
+```shell
+export ZEN_API_KEY=your-api-key
+python tests/pentest.py
+```
+
+Run pentest against both ZenGuard AI and ChatGPT:
+```shell
+export ZEN_API_KEY=your-api-key
+export OPENAI_API_KEY=your-openai-api-key
+python tests/pentest.py
+```
+
+
+Note that we always are running the pentest against the most up-to-date model. Currently, `gpt-4-0125-preview`
+
+
+
 # Support and Contact
 
 [Book a Demo](https://calendly.com/galym-u) or just shoot us an email to hello@zenguard.ai.

pyproject.toml

@@ -1,18 +1,15 @@
 [tool.poetry]
 name = "zenguard"
-version = "0.1.3"
+version = "0.1.7"
 description = "Plug-and-play production grade security for GenAI applications"
 authors = ["ZenGuard Team <hello@zenguard.ai>"]
 license = "MIT"
 readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "^3.9"
-httpx = "^0.27.0"
+httpx = ">=0.24.0,<1.0.0"
 tqdm = "^4.66.2"
-
-
-[tool.poetry.group.pentest.dependencies]
 openai = "^1.14.2"
 rapidfuzz = "^3.7.0"
 pandas = "^2.2.1"

tests/pentest.py

@@ -1,11 +1,28 @@
 import os
 
-from zenguard import Credentials, Detector, Endpoint, ZenGuard, ZenGuardConfig
+from zenguard import (
+    Credentials,
+    Detector,
+    Endpoint,
+    SupportedLLMs,
+    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")
+    openai_api_key = os.environ.get("OPENAI_API_KEY")
+    if not openai_api_key:
+        print("Pentesting ZenGuard only (OpenAI API key not set)")
+
     config = ZenGuardConfig(credentials=Credentials(api_key=api_key))
+    if openai_api_key:
+        config.credentials.llm_api_key = openai_api_key
+        config.llm = SupportedLLMs.CHATGPT
     zenguard = ZenGuard(config=config)
-    zenguard.pentest(endpoint=Endpoint.ZENGUARD, detector=Detector.PROMPT_INJECTION)
+    zenguard.pentest(endpoint=Endpoint.ZENGUARD, detector=Detector.PROMPT_INJECTION)
+
+    if openai_api_key:
+        zenguard.pentest(endpoint=Endpoint.OPENAI, detector=Detector.PROMPT_INJECTION)

zenguard/__init__.py

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

zenguard/pentest/prompt_injections/run.py

@@ -1,16 +1,15 @@
-import openai
 from tqdm import tqdm
 
 
-def run_prompts_api(built_prompts, quiet=False, dry_run=False):
+def run_prompts_api(built_prompts, client, quiet=False, dry_run=False):
     if not quiet:
         built_prompts = tqdm(built_prompts)
     for prompt in built_prompts:
         if dry_run:
             api_result = _get_mocked_api_response()
             prompt["result"] = api_result["choices"][0]["message"]["content"]
         else:
-            api_result = _prompt_model_api(prompt)
+            api_result = _prompt_model_api(prompt, client)
             prompt["result"] = api_result.choices[0].message.content
             print(prompt["result"])
 
@@ -32,7 +31,7 @@ def _get_mocked_api_response():
     }
 
 
-def _prompt_model_api(prompt, use_stop=False):
+def _prompt_model_api(prompt, client, use_stop=False):
     prompt_settings = prompt["settings"]
 
     api_prompt_string = prompt["prompt"]
@@ -45,7 +44,7 @@ def _prompt_model_api(prompt, use_stop=False):
 
     messages = [{"role": "user", "content": api_prompt_string}]
 
-    response = openai.chat.completions.create(
+    response = client.chat.completions.create(
         model=api_config_model,
         messages=messages,
         temperature=api_config_temperature,

zenguard/zenguard.py

@@ -4,21 +4,35 @@
 
 from dataclasses import dataclass
 from enum import Enum
+from typing import Optional
 
 import httpx
+from openai import OpenAI
 from tqdm import tqdm
 
-from zenguard.pentest.prompt_injections import config, prompting, scoring, visualization
+from zenguard.pentest.prompt_injections import (
+    config,
+    prompting,
+    run,
+    scoring,
+    visualization,
+)
+
+
+class SupportedLLMs:
+    CHATGPT = "chatgpt"
 
 
 @dataclass
 class Credentials:
     api_key: str
+    llm_api_key: Optional[str] = None
 
 
 @dataclass
 class ZenGuardConfig:
     credentials: Credentials
+    llm: Optional[SupportedLLMs] = None
 
 
 class Detector(Enum):
@@ -29,11 +43,11 @@ class Detector(Enum):
     KEYWORDS = "v1/detect/keywords"
     SECRETS = "v1/detect/secrets"
 
+
 class Endpoint(Enum):
     ZENGUARD = "zenguard"
     OPENAI = "openai"
 
-
 class ZenGuard:
     """
     ZenGuard is a class that represents the ZenGuard object.
@@ -47,6 +61,14 @@ def __init__(
         self._api_key = config.credentials.api_key
         self._backend = "https://api.zenguard.ai/"
 
+        self._llm_client = None
+        if config.llm == SupportedLLMs.CHATGPT:
+            self._llm_client = OpenAI(
+                api_key=config.credentials.llm_api_key,
+            )
+        elif config.llm is not None:
+            raise ValueError(f"LLM {config.llm} is not supported")
+
     def detect(self, detectors: list[Detector], prompt: str):
         if len(detectors) == 0:
             return {"error": "No detectors were provided"}
@@ -83,6 +105,9 @@ def pentest(self, endpoint: Endpoint, detector: Detector = None):
                 detector == Detector.PROMPT_INJECTION
             ), "Only prompt injection pentesting is currently supported"
             self._attack_zenguard(Detector.PROMPT_INJECTION, attack_prompts)
+        elif endpoint == Endpoint.OPENAI:
+            print("Running attack on OpenAI endpoint")
+            run.run_prompts_api(attack_prompts, self._llm_client)
 
         scoring.score_attacks(attack_prompts)
         df = visualization.build_dataframe(attack_prompts)