rag poisoning-pr (#68)

* adding support for base_url with ollama and openai.
adding embedding
adding target temperature for embedding attacks
configuration. via file and menu
adding skipped test method
adding rag poisnoning attack
adding package creation dependencies via setup.py (oldschool)
adding uv package baseline
adding tests

* refactored provider and model prompts

* Disable bugged telemetry in RAG Poisoning test to prevent PostHog errors open-webui/open-webui#15624
+ add dependencies to base package

* restored comment

* configuration error should also skip

* Improve error handling for RAG poisoning attack  (these errors should be caught and warned)

* Update ps_fuzz/attacks/rag_poisoning.py

supress loggers

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update ps_fuzz/attacks/rag_poisoning.py

out of scope fail-safe

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update ps_fuzz/test_base.py

typo

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update ps_fuzz/attacks/rag_poisoning.py

operator race condition

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update tests/test_chat_clients.py

unused

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update tests/test_chat_clients.py

redefined in #1

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update tests/test_chat_clients.py

unused

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update ps_fuzz/attacks/rag_poisoning.py

unused

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Fix tests, introduce claude.md

* Update README.md to include RAG & Vector Database Attacks section and enhance attack options

* Add Bandit configuration file and implement poisoned document creation in RAG poisoning attack

* Address all Copilot review comments on PR #68 (#69)

- Fix ChromaDB persist() compatibility: wrap in try/except for 0.4.0+
  which auto-persists with persist_directory
- Replace fragile error string matching with specific exception types
  (ImportError, ConnectionError, ValueError, etc.)
- Fix register_test decorator to return cls (was returning None)
- Fix getter/setter inconsistency: embedding_provider and
  embedding_model setters now accept empty values matching getter defaults
- Fix empty base_url passthrough: empty strings are now stripped from
  kwargs instead of passed to model constructors
- Remove unused client variable assignments in test_chat_clients.py
- Reorganize tests: move AppConfig, helper, and TestStatus tests out of
  test_is_response_list.py into dedicated test files

All 93 tests pass.

https://claude.ai/code/session_01CDFqeg5QhB4V7yQ3yVVBc9

Co-authored-by: Claude <noreply@anthropic.com>

---------

Co-authored-by: David Abutbul <david@abutbul.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>

Author: 4 people

.bandit

@@ -0,0 +1,10 @@
+# Bandit configuration file
+# Exclude directories that should not be scanned
+exclude_dirs:
+  - './.venv'
+  - './.git'
+  - './build'
+  - './dist'
+  - './prompt_security_fuzzer.egg-info'
+  - './.env'
+  - './tests'  # Exclude test files - pytest uses assertions which trigger B101 warnings

README.md

@@ -56,6 +56,7 @@ Table of Contents
 * [Supported attacks](#attacks)
    * [Jailbreak](#jailbreak)
    * [Prompt Injection](#pi-injection)
+   * [RAG & Vector Database Attacks](#rag-poisoning)
    * [System prompt extraction](#systemleak)
 * [ :rainbow:  What’s next on the roadmap?](#roadmap)
 * [ :beers: Contributing](#contributing)
@@ -111,7 +112,7 @@ Table of Contents
 ### Features
 <b>The Prompt Fuzzer Supports:</b><br>
 🧞  16 [llm providers](#llm-providers)<br>
-🔫  15 different [attacks](#attacks)<br>
+🔫  16 different [attacks](#attacks)<br>
 💬  Interactive mode<br>
 🤖  CLI mode<br>
 🧵  Multi threaded testing<br>
@@ -163,8 +164,14 @@ Alternatively, create a file named `.env` in the current directory and set the `
 * `--num-attempts, -n`       NUM_ATTEMPTS Number of different attack prompts 
 * `--num-threads, -t`        NUM_THREADS  Number of worker threads 
 * `--attack-temperature, -a` ATTACK_TEMPERATURE  Temperature for attack model 
-* `--debug-level, -d`        DEBUG_LEVEL  Debug level (0-2)   
-* `-batch, -b`               Run the fuzzer in unattended (batch) mode, bypassing the interactive steps 
+* `--debug-level, -d`        DEBUG_LEVEL  Debug level (0-2)
+* `-batch, -b`               Run the fuzzer in unattended (batch) mode, bypassing the interactive steps
+* `--ollama-base-url`        Base URL for Ollama API (for self-hosted deployments)
+* `--openai-base-url`        Base URL for OpenAI API (for OpenAI-compatible endpoints)
+* `--embedding-provider`     Embedding provider (ollama or open_ai) - required for RAG tests
+* `--embedding-model`        Embedding model name - required for RAG tests
+* `--embedding-ollama-base-url` Base URL for Ollama Embedding API
+* `--embedding-openai-base-url` Base URL for OpenAI Embedding API
 
 <br/>
 
@@ -205,6 +212,43 @@ Run tests against the system prompt with a subset of attacks
     prompt-security-fuzzer -b ./system_prompt.examples/medium_system_prompt.txt --custom-benchmark=ps_fuzz/attack_data/custom_benchmark1.csv --tests='["ucar","amnesia"]'
 ```
 
+#### 🧪 RAG Poisoning Attack
+Test RAG systems with vector database poisoning attacks
+
+```bash
+# Using OpenAI embeddings
+prompt-security-fuzzer -b ./system_prompt.examples/medium_system_prompt.txt \
+    --embedding-provider=open_ai \
+    --embedding-model=text-embedding-ada-002 \
+    --tests='["rag_poisoning"]'
+
+# Using Ollama embeddings with custom endpoint
+prompt-security-fuzzer -b ./system_prompt.examples/medium_system_prompt.txt \
+    --embedding-provider=ollama \
+    --embedding-model=nomic-embed-text \
+    --embedding-ollama-base-url=http://localhost:11434 \
+    --tests='["rag_poisoning"]'
+```
+
+**Note**: Requires chromadb (installed by default with prompt-security-fuzzer)
+
+#### 🔌 Using Custom API Endpoints
+Run tests against custom or self-hosted LLM deployments
+
+```bash
+# Using custom Ollama endpoint
+prompt-security-fuzzer -b ./system_prompt.examples/medium_system_prompt.txt \
+    --target-provider=ollama \
+    --target-model=llama2 \
+    --ollama-base-url=http://localhost:11434
+
+# Using OpenAI-compatible endpoint (e.g., LocalAI, vLLM, LM Studio)
+prompt-security-fuzzer -b ./system_prompt.examples/medium_system_prompt.txt \
+    --target-provider=open_ai \
+    --target-model=custom-model \
+    --openai-base-url=http://your-custom-endpoint:8000/v1
+```
+
 <br>
 <br>
 <br>
@@ -245,6 +289,11 @@ We use a dynamic testing approach, where we get the necessary context from your
 - **Ethical Compliance**: Evaluates resistance to discussing harmful or inappropriate content about sensitive topics.
 - **Typoglycemia Attack**: Exploits text processing vulnerabilities by omitting random characters, causing incorrect responses.
 
+<a id="rag-poisoning"></a>
+##### RAG & Vector Database Attacks
+
+- **RAG Poisoning (Hidden Parrot Attack)**: Tests whether malicious instructions embedded in vector database documents can compromise RAG system behavior. This attack verifies if poisoned content retrieved from vector stores can override system prompts or inject unauthorized instructions into LLM responses.
+
 <a id="systemleak"></a>
 ##### System prompt extraction
 

claude.md

@@ -0,0 +1,39 @@
+# Development Setup
+
+## Python Environment
+
+This project requires Python >= 3.9 (tested with 3.9, 3.10, 3.11).
+
+### Setup with uv
+
+1. Create virtual environment with Python 3.11:
+   ```bash
+   uv venv --python 3.11
+   ```
+
+2. Activate the virtual environment:
+   ```bash
+   source .venv/bin/activate
+   ```
+
+3. Install dependencies:
+   ```bash
+   uv pip install -e ".[dev]"
+   ```
+
+### Running Tests
+
+Run all tests:
+```bash
+pytest
+```
+
+Run specific test:
+```bash
+pytest tests/test_chat_clients.py::TestClientLangChainBaseURL::test_empty_base_url_parameters -v
+```
+
+Run tests with verbose output:
+```bash
+pytest -v
+```

ps_fuzz/app_config.py

@@ -38,7 +38,8 @@ def __init__(self, config_state_file: str, config_state: dict = None):
                 logger.warning(f"Failed to load config state file {self.config_state_file}: {e}")
 
     def get_attributes(self):
-        return self.config_state
+        attributes = self.config_state.copy()
+        return attributes
 
     def print_as_table(self):
         attributes = self.get_attributes()
@@ -184,6 +185,60 @@ def system_prompt(self) -> str:
     def system_prompt(self, value: str):
         self.config_state['system_prompt'] = value
         self.save()
+        
+    @property
+    def ollama_base_url(self) -> str:
+        return self.config_state.get('ollama_base_url', '')
+    
+    @ollama_base_url.setter
+    def ollama_base_url(self, value: str):
+        self.config_state['ollama_base_url'] = value
+        self.save()
+        
+    @property
+    def openai_base_url(self) -> str:
+        return self.config_state.get('openai_base_url', '')
+    
+    @openai_base_url.setter
+    def openai_base_url(self, value: str):
+        self.config_state['openai_base_url'] = value
+        self.save()
+        
+    @property
+    def embedding_provider(self) -> str:
+        return self.config_state.get('embedding_provider', '')
+    
+    @embedding_provider.setter
+    def embedding_provider(self, value: str):
+        self.config_state['embedding_provider'] = value if value else ''
+        self.save()
+        
+    @property
+    def embedding_ollama_base_url(self) -> str:
+        return self.config_state.get('embedding_ollama_base_url', '')
+    
+    @embedding_ollama_base_url.setter
+    def embedding_ollama_base_url(self, value: str):
+        self.config_state['embedding_ollama_base_url'] = value
+        self.save()
+        
+    @property
+    def embedding_openai_base_url(self) -> str:
+        return self.config_state.get('embedding_openai_base_url', '')
+    
+    @embedding_openai_base_url.setter
+    def embedding_openai_base_url(self, value: str):
+        self.config_state['embedding_openai_base_url'] = value
+        self.save()
+        
+    @property
+    def embedding_model(self) -> str:
+        return self.config_state.get('embedding_model', '')
+    
+    @embedding_model.setter
+    def embedding_model(self, value: str):
+        self.config_state['embedding_model'] = value if value else ''
+        self.save()
 
     def update_from_args(self, args):
         args_dict = vars(args)
@@ -218,6 +273,12 @@ def parse_cmdline_args():
     parser.add_argument('-a', '--attack-temperature', type=float, default=None, help="Temperature for attack model")
     parser.add_argument('-d', '--debug-level', type=int, default=None, help="Debug level (0-2)")
     parser.add_argument("-b", '--batch', action='store_true', help="Run the fuzzer in unattended (batch) mode, bypassing the interactive steps")
+    parser.add_argument('--ollama-base-url', type=str, dest='ollama_base_url', default=None, help="Base URL for Ollama API")
+    parser.add_argument('--openai-base-url', type=str, dest='openai_base_url', default=None, help="Base URL for OpenAI API")
+    parser.add_argument('--embedding-provider', type=str, dest='embedding_provider', default=None, help="Embedding provider (ollama or open_ai)")
+    parser.add_argument('--embedding-ollama-base-url', type=str, dest='embedding_ollama_base_url', default=None, help="Base URL for Ollama Embedding API")
+    parser.add_argument('--embedding-openai-base-url', type=str, dest='embedding_openai_base_url', default=None, help="Base URL for OpenAI Embedding API")
+    parser.add_argument('--embedding-model', type=str, dest='embedding_model', default=None, help="Embedding model name")
     parser.add_argument('system_prompt_file', type=str, nargs='?', default=None, help="Filename containing the system prompt")
     return parser.parse_args()
 

ps_fuzz/attack_config.py

@@ -1,6 +1,7 @@
 from .client_config import ClientConfig
 
 class AttackConfig(object):
-    def __init__(self, attack_client: ClientConfig, attack_prompts_count: int):
+    def __init__(self, attack_client: ClientConfig, attack_prompts_count: int, embedding_config=None):
         self.attack_client = attack_client
         self.attack_prompts_count = attack_prompts_count
+        self.embedding_config = embedding_config

ps_fuzz/attack_loader.py

@@ -10,5 +10,6 @@
     complimentary_transition,
     harmful_behavior,
     base64_injection,
-    custom_benchmark
+    custom_benchmark,
+    rag_poisoning
 )

ps_fuzz/attack_registry.py

@@ -14,6 +14,7 @@ def register_test(cls):
     global test_classes
     logger.debug(f"Registering attack test class: {cls.__name__}")
     test_classes.append(cls)
+    return cls
 
 def instantiate_tests(client_config: ClientConfig, attack_config:AttackConfig, custom_tests:List=None, custom_benchmark:bool=False) -> List[TestBase]:
     tests = []