mirror of
https://github.com/Shiva108/ai-llm-red-team-handbook.git
synced 2026-07-11 14:56:32 +02:00
shiva108
56 lines
2.0 KiB
Python
Executable File
56 lines
2.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
RAG-Focused Exploitation Workflow
|
|
|
|
Specialized workflow for attacking RAG (Retrieval Augmented Generation) systems.
|
|
Combines vector database poisoning, retrieval manipulation, and indirect injection.
|
|
|
|
Usage:
|
|
python3 workflows/rag_exploitation.py --target https://api.example.com --vector-db chromadb
|
|
"""
|
|
|
|
import argparse
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.append(str(Path(__file__).parent.parent))
|
|
|
|
def main():
|
|
"""RAG exploitation workflow."""
|
|
parser = argparse.ArgumentParser(description='RAG-focused exploitation workflow')
|
|
parser.add_argument('--target', required=True, help='Target RAG API URL')
|
|
parser.add_argument('--vector-db', choices=['chromadb', 'faiss', 'pinecone'], help='Vector database type')
|
|
parser.add_argument('--poison-docs', help='Documents for poisoning attack')
|
|
parser.add_argument('--output', '-o', help='Output report file')
|
|
parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
|
|
|
|
args = parser.parse_args()
|
|
|
|
print(f"RAG Exploitation Workflow")
|
|
print(f"Target: {args.target}")
|
|
print(f"Vector DB: {args.vector_db}")
|
|
print(f"\nPhase 1: RAG Architecture Reconnaissance")
|
|
print(" - Identifying retrieval endpoints")
|
|
print(" - Analyzing embedding model")
|
|
print(" - Mapping vector database")
|
|
|
|
print(f"\nPhase 2: Vector Database Poisoning")
|
|
print(" - Crafting poisoned documents")
|
|
print(" - Injecting malicious embeddings")
|
|
print(" - Testing retrieval manipulation")
|
|
|
|
print(f"\nPhase 3: Indirect Injection via RAG")
|
|
print(" - Embedding hidden instructions")
|
|
print(" - Testing context injection")
|
|
print(" - Validating payload execution")
|
|
|
|
print(f"\nPhase 4: Data Extraction")
|
|
print(" - Extracting indexed documents")
|
|
print(" - Leaking embedding vectors")
|
|
print(" - Recovering training data")
|
|
|
|
print("\n[!] This is a template workflow - implement actual attack logic")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|