""" Phase 7 Stress Testing Runner. Run adversarial tests against the RAG pipeline to identify vulnerabilities. """ import argparse import json import logging import sys from pathlib import Path from src.stress_testing.adapter import ( AttackCategory, StressTestingAdapter, ) from src.stress_testing.bias_probing import BiasProbingTester from src.stress_testing.information_evasion import InformationEvasionTester from src.stress_testing.prompt_injection import PromptInjectionTester logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", ) logger = logging.getLogger(__name__) def get_all_test_cases() -> list: """Get all stress test cases from all categories.""" all_cases = [] all_cases.extend(PromptInjectionTester.get_test_cases()) all_cases.extend(InformationEvasionTester.get_test_cases()) all_cases.extend(BiasProbingTester.get_test_cases()) return all_cases def run_stress_tests( output_path: Path | None = None, verbose: bool = False, use_live: bool = False, test_cases: list | None = None, ) -> dict: """ Run all stress tests and generate a report. Args: output_path: Optional path to save JSON report. verbose: Print detailed results. use_live: If True, run against live ReasoningPipeline. test_cases: Optional list of test cases to run. If None, uses all cases. Returns: Dictionary with test report. """ print("Initializing stress testing adapter...") if use_live: print("Using LIVE Reasoning Pipeline - tests will run against actual system") print("NOTE: Make sure your LLM service (Ollama/OpenRouter) is running!") print() adapter = StressTestingAdapter(use_live_pipeline=True) else: print("Using SIMULATION mode - no live pipeline connected") print("Use --live flag to test against the actual reasoning engine") print() adapter = StressTestingAdapter() all_cases = test_cases if test_cases is not None else get_all_test_cases() print(f"Running {len(all_cases)} stress tests...") print() # Run tests with progress indicator report = adapter.run_tests(all_cases) print("\n" + "=" * 60) print("STRESS TESTING REPORT") print("=" * 60) print() print(f"Mode: {'LIVE' if use_live else 'SIMULATION'}") print(f"Total Tests: {report.total_tests}") print(f"Passed Defenses: {report.passed_defenses}") print(f"Failed Defenses: {report.failed_defenses}") print(f"Defense Rate: {report.defense_rate:.1%}") print() if verbose: print("Individual Results:") for result in report.results: status = "PASS" if result.defense_triggered else "FAIL" latency = f"{result.latency_ms:.0f}ms" if result.latency_ms else "N/A" print(f" [{status}] {result.attack_name} ({result.category.value}) - {latency}") if result.notes: print(f" Note: {result.notes}") print() print("By Category:") for category in AttackCategory: cat_results = [r for r in report.results if r.category == category] cat_passed = sum(1 for r in cat_results if r.defense_triggered) cat_total = len(cat_results) cat_rate = cat_passed / cat_total if cat_total > 0 else 0 print(f" {category.value}: {cat_passed}/{cat_total} ({cat_rate:.0%})") if output_path: with output_path.open("w") as f: json.dump(report.to_dict(), f, indent=2) print(f"\nReport saved to: {output_path}") return report.to_dict() def run_by_category( category: AttackCategory, output_path: Path | None = None, verbose: bool = False, use_live: bool = False, ) -> dict: """ Run stress tests for a specific category. Args: category: The attack category to test. output_path: Optional path to save JSON report. verbose: Print detailed results. use_live: If True, run against live pipeline. Returns: Dictionary with test report. """ print(f"Running {category.value} tests...") adapter = StressTestingAdapter(use_live_pipeline=True) if use_live else StressTestingAdapter() all_cases = get_all_test_cases() filtered_cases = [c for c in all_cases if c.category == category] print(f"Running {len(filtered_cases)} tests in {category.value}...") report = adapter.run_tests(filtered_cases) print() print(f"Defense Rate: {report.defense_rate:.1%}") if verbose: print("\nResults:") for result in report.results: status = "PASS" if result.defense_triggered else "FAIL" print(f" [{status}] {result.attack_name}") if output_path: with output_path.open("w") as f: json.dump(report.to_dict(), f, indent=2) print(f"Report saved to: {output_path}") return report.to_dict() def main() -> None: """Main entry point.""" parser = argparse.ArgumentParser(description="Phase 7 Stress Testing - Red Teaming for RAG Pipeline") parser.add_argument( "--category", choices=["prompt_injection", "information_evasion", "bias_probing"], help="Run tests for a specific category", ) parser.add_argument( "--output", type=Path, help="Save report to JSON file", ) parser.add_argument( "--verbose", "-v", action="store_true", help="Print detailed results", ) parser.add_argument( "--live", action="store_true", help=( "Run against live ReasoningPipeline instead of simulation mode. " "Requires LLM service running (Ollama or OpenRouter)." ), ) parser.add_argument( "--limit", type=int, help="Limit number of tests to run (useful for quick testing)", ) args = parser.parse_args() try: # Get test cases all_cases = get_all_test_cases() # Apply limit if specified limited_cases = all_cases[: args.limit] if args.limit else None if args.limit: print(f"Limiting to first {args.limit} tests") if args.category: category = AttackCategory(args.category) run_by_category(category, args.output, verbose=args.verbose, use_live=args.live) else: run_stress_tests( args.output, verbose=args.verbose, use_live=args.live, test_cases=limited_cases, ) except RuntimeError as e: print(f"\nERROR: {e}", file=sys.stderr) print("\nTo fix:", file=sys.stderr) print("1. If using Ollama: Ensure 'ollama serve' is running", file=sys.stderr) print("2. If using OpenRouter: Set OPENROUTER_API_KEY", file=sys.stderr) print("3. Check .env configuration", file=sys.stderr) sys.exit(1) except Exception as e: logger.exception("Unexpected error") print(f"\nERROR: {e}", file=sys.stderr) sys.exit(1) if __name__ == "__main__": main()