mirror of
https://github.com/Shiva108/ai-llm-red-team-handbook.git
synced 2026-08-27 21:30:27 +02:00
feat: Add script to extract code blocks from markdown files and generate a JSON catalog.
This commit is contained in:
shiva108
Executable
+1
@@ -0,0 +1 @@
|
||||
"""Workflows module for AI LLM Red Teaming."""
|
||||
Executable
+190
@@ -0,0 +1,190 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Full LLM Security Assessment Workflow
|
||||
|
||||
This workflow orchestrates a comprehensive security assessment combining
|
||||
reconnaissance, prompt injection, data extraction, and exploitation techniques.
|
||||
|
||||
Usage:
|
||||
python3 workflows/full_assessment.py --target https://api.example.com --output report.json
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import json
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
|
||||
# Add parent to path for imports
|
||||
sys.path.append(str(Path(__file__).parent.parent))
|
||||
|
||||
class AssessmentOrchestrator:
|
||||
"""Orchestrates a full LLM security assessment."""
|
||||
|
||||
def __init__(self, target, output_file=None, verbose=False):
|
||||
self.target = target
|
||||
self.output_file = output_file
|
||||
self.verbose = verbose
|
||||
self.results = {
|
||||
'target': target,
|
||||
'timestamp': datetime.now().isoformat(),
|
||||
'phases': {}
|
||||
}
|
||||
|
||||
def log(self, message):
|
||||
"""Log message if verbose."""
|
||||
if self.verbose:
|
||||
print(f"[*] {message}")
|
||||
|
||||
def run_phase(self, phase_name, description, scripts):
|
||||
"""Run a phase of the assessment."""
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Phase: {phase_name}")
|
||||
print(f"Description: {description}")
|
||||
print(f"{'='*60}\n")
|
||||
|
||||
phase_results = []
|
||||
|
||||
for script_name, script_desc in scripts:
|
||||
self.log(f"Running: {script_name}")
|
||||
self.log(f"Purpose: {script_desc}")
|
||||
|
||||
# TODO: Actually execute the scripts
|
||||
# For now, just record them
|
||||
phase_results.append({
|
||||
'script': script_name,
|
||||
'description': script_desc,
|
||||
'status': 'planned'
|
||||
})
|
||||
|
||||
self.results['phases'][phase_name] = phase_results
|
||||
return phase_results
|
||||
|
||||
def phase_1_reconnaissance(self):
|
||||
"""Phase 1: Reconnaissance and fingerprinting."""
|
||||
scripts = [
|
||||
('reconnaissance/chapter_31_ai_system_reconnaissance_01_reconnaissance.py',
|
||||
'LLM system fingerprinting'),
|
||||
('reconnaissance/chapter_31_ai_system_reconnaissance_02_reconnaissance.py',
|
||||
'API discovery and enumeration'),
|
||||
]
|
||||
return self.run_phase('Reconnaissance',
|
||||
'Identify LLM type, version, and architecture',
|
||||
scripts)
|
||||
|
||||
def phase_2_prompt_injection(self):
|
||||
"""Phase 2: Prompt injection testing."""
|
||||
scripts = [
|
||||
('prompt_injection/chapter_14_prompt_injection_01_prompt_injection.py',
|
||||
'Basic prompt injection'),
|
||||
('prompt_injection/chapter_14_prompt_injection_02_prompt_injection.py',
|
||||
'Context overflow attacks'),
|
||||
('prompt_injection/chapter_14_prompt_injection_03_prompt_injection.py',
|
||||
'System prompt leakage'),
|
||||
]
|
||||
return self.run_phase('Prompt Injection',
|
||||
'Test prompt injection vulnerabilities',
|
||||
scripts)
|
||||
|
||||
def phase_3_data_extraction(self):
|
||||
"""Phase 3: Data extraction attempts."""
|
||||
scripts = [
|
||||
('data_extraction/chapter_15_data_leakage_and_extraction_01_data_extraction.py',
|
||||
'PII extraction'),
|
||||
('data_extraction/chapter_15_data_leakage_and_extraction_02_data_extraction.py',
|
||||
'Training data extraction'),
|
||||
('data_extraction/chapter_15_data_leakage_and_extraction_03_data_extraction.py',
|
||||
'Memory dump attempts'),
|
||||
]
|
||||
return self.run_phase('Data Extraction',
|
||||
'Attempt to extract sensitive data',
|
||||
scripts)
|
||||
|
||||
def phase_4_jailbreak(self):
|
||||
"""Phase 4: Jailbreak attempts."""
|
||||
scripts = [
|
||||
('jailbreak/chapter_16_jailbreaks_and_bypass_techniques_01_jailbreak.py',
|
||||
'Character roleplay bypass'),
|
||||
('jailbreak/chapter_16_jailbreaks_and_bypass_techniques_02_jailbreak.py',
|
||||
'DAN techniques'),
|
||||
]
|
||||
return self.run_phase('Jailbreak Testing',
|
||||
'Test guardrail bypasses',
|
||||
scripts)
|
||||
|
||||
def phase_5_plugin_exploitation(self):
|
||||
"""Phase 5: Plugin and API exploitation."""
|
||||
scripts = [
|
||||
('plugin_exploitation/chapter_17_01_fundamentals_and_architecture_01_plugin_exploitation.py',
|
||||
'Plugin enumeration'),
|
||||
('plugin_exploitation/chapter_17_02_api_authentication_and_authorization_01_plugin_exploitation.py',
|
||||
'Authentication bypass'),
|
||||
('plugin_exploitation/chapter_17_04_api_exploitation_and_function_calling_01_plugin_exploitation.py',
|
||||
'Command injection'),
|
||||
]
|
||||
return self.run_phase('Plugin Exploitation',
|
||||
'Test plugin and API vulnerabilities',
|
||||
scripts)
|
||||
|
||||
def phase_6_reporting(self):
|
||||
"""Phase 6: Generate final report."""
|
||||
print(f"\n{'='*60}")
|
||||
print("Generating Assessment Report")
|
||||
print(f"{'='*60}\n")
|
||||
|
||||
if self.output_file:
|
||||
Path(self.output_file).write_text(json.dumps(self.results, indent=2))
|
||||
print(f"Report saved to: {self.output_file}")
|
||||
else:
|
||||
print(json.dumps(self.results, indent=2))
|
||||
|
||||
def run_full_assessment(self):
|
||||
"""Run the complete assessment workflow."""
|
||||
print(f"\n{'#'*60}")
|
||||
print(f"# AI LLM Security Assessment")
|
||||
print(f"# Target: {self.target}")
|
||||
print(f"# Time: {self.results['timestamp']}")
|
||||
print(f"{'#'*60}\n")
|
||||
|
||||
# Run all phases
|
||||
self.phase_1_reconnaissance()
|
||||
self.phase_2_prompt_injection()
|
||||
self.phase_3_data_extraction()
|
||||
self.phase_4_jailbreak()
|
||||
self.phase_5_plugin_exploitation()
|
||||
self.phase_6_reporting()
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
print("Assessment Complete")
|
||||
print(f"{'='*60}\n")
|
||||
|
||||
def main():
|
||||
"""Command-line interface."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description='Full LLM Security Assessment Workflow',
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
# Run full assessment
|
||||
python3 workflows/full_assessment.py --target https://api.example.com
|
||||
|
||||
# Save results to file
|
||||
python3 workflows/full_assessment.py --target https://api.example.com --output report.json
|
||||
|
||||
# Verbose mode
|
||||
python3 workflows/full_assessment.py --target https://api.example.com --verbose
|
||||
"""
|
||||
)
|
||||
|
||||
parser.add_argument('--target', required=True, help='Target LLM API URL')
|
||||
parser.add_argument('--output', '-o', help='Output file for results (JSON)')
|
||||
parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Run assessment
|
||||
orchestrator = AssessmentOrchestrator(args.target, args.output, args.verbose)
|
||||
orchestrator.run_full_assessment()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+58
@@ -0,0 +1,58 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Plugin Penetration Testing Workflow
|
||||
|
||||
Focused workflow for testing plugin and API vulnerabilities in LLM systems.
|
||||
Covers authentication, authorization, command injection, and function hijacking.
|
||||
|
||||
Usage:
|
||||
python3 workflows/plugin_pentest.py --target https://api.example.com --plugins weather,calculator
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.append(str(Path(__file__).parent.parent))
|
||||
|
||||
def main():
|
||||
"""Plugin penetration testing workflow."""
|
||||
parser = argparse.ArgumentParser(description='Plugin penetration testing workflow')
|
||||
parser.add_argument('--target', required=True, help='Target LLM API URL')
|
||||
parser.add_argument('--plugins', help='Comma-separated list of plugins to test')
|
||||
parser.add_argument('--auth-token', help='Authentication token (if required)')
|
||||
parser.add_argument('--output', '-o', help='Output report file')
|
||||
parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
plugins = args.plugins.split(',') if args.plugins else []
|
||||
|
||||
print(f"Plugin Penetration Testing Workflow")
|
||||
print(f"Target: {args.target}")
|
||||
print(f"Plugins: {', '.join(plugins)}")
|
||||
|
||||
print(f"\nPhase 1: Plugin Discovery")
|
||||
print(" - Enumerating available plugins")
|
||||
print(" - Analyzing plugin manifests")
|
||||
print(" - Mapping function signatures")
|
||||
|
||||
print(f"\nPhase 2: Authentication Testing")
|
||||
print(" - Testing API key requirements")
|
||||
print(" - Attempting token bypass")
|
||||
print(" - Checking authorization boundaries")
|
||||
|
||||
print(f"\nPhase 3: Function Calling Exploitation")
|
||||
print(" - Testing command injection in plugins")
|
||||
print(" - Attempting function hijacking")
|
||||
print(" - Exploiting type confusion")
|
||||
|
||||
print(f"\nPhase 4: Third-Party Integration Testing")
|
||||
print(" - Analyzing external API calls")
|
||||
print(" - Testing SSRF vulnerabilities")
|
||||
print(" - Checking credential leakage")
|
||||
|
||||
print("\n[!] This is a template workflow - implement actual testing logic")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+55
@@ -0,0 +1,55 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user