diff --git a/tools/prompt_injection_tester/pit/commands/scan.py b/tools/prompt_injection_tester/pit/commands/scan.py index 582058b..cd58a2f 100644 --- a/tools/prompt_injection_tester/pit/commands/scan.py +++ b/tools/prompt_injection_tester/pit/commands/scan.py @@ -20,6 +20,8 @@ from pit.ui.console import console from pit.ui.display import print_banner, print_success, print_error, print_warning, print_info from pit.ui.progress import create_progress_bar, create_spinner from pit.ui.tables import create_results_table, print_summary_panel +from pit.orchestrator import WorkflowOrchestrator +from pit.orchestrator.workflow import WorkflowOrchestrator # Create sub-app for scan commands @@ -174,68 +176,67 @@ async def _run_auto_scan( """ print_info("Phase 1: Discovery & Reconnaissance") - with create_spinner() as progress: - task = progress.add_task("Discovering endpoint...", total=None) + # Parse patterns + pattern_list = patterns.split(",") if patterns else None - # Simulate discovery - await asyncio.sleep(1) - discovered_model = model or "llama3:latest" - - progress.update(task, description=f"✓ Found model: {discovered_model}") - - print_success(f"Discovered model: {discovered_model}") - - print_info("Phase 2: Loading attack patterns") - - pattern_list = patterns.split(",") if patterns else [ - "direct_instruction_override", - "direct_role_authority", - "direct_persona_shift", - ] - - print_success(f"Loaded {len(pattern_list)} attack patterns") - - print_info("Phase 3: Executing attacks") - - results = [] - with create_progress_bar() as progress: - task = progress.add_task( - "Testing patterns...", - total=len(pattern_list), - ) - - for pattern in pattern_list: - # Simulate attack execution - await asyncio.sleep(0.5) - - result = { - "pattern": pattern, - "success": True, # Placeholder - "confidence": 0.85, - "details": "Pattern executed successfully", - } - results.append(result) - - progress.update(task, advance=1) - - print_success(f"Completed {len(results)} attack tests") - - # Display results - console.print() - table = create_results_table(results) - console.print(table) - - # Display summary - console.print() - successful = sum(1 for r in results if r["success"]) - failed = len(results) - successful - print_summary_panel( - total=len(results), - successful=successful, - failed=failed, - duration=2.5, + # Create orchestrator + orchestrator = WorkflowOrchestrator( + target_url=target, + model=model, + verbose=verbose, ) + try: + # Run workflow with progress tracking + with create_progress_bar() as progress: + task = progress.add_task( + "Running security assessment...", + total=100, + ) + + # Execute workflow + workflow_results = await orchestrator.run_auto_workflow( + patterns=pattern_list, + progress=progress, + progress_task=task, + ) + + progress.update(task, completed=100) + + # Check for errors + if workflow_results.get("errors"): + for error in workflow_results["errors"]: + print_error(f"Error: {error}") + + if not workflow_results.get("success"): + return + + # Display discovered model + discovered_model = workflow_results.get("model", "unknown") + print_success(f"Discovered model: {discovered_model}") + + # Display results + results = workflow_results.get("tests", []) + if results: + console.print() + table = create_results_table(results) + console.print(table) + + # Display summary + console.print() + summary = workflow_results.get("summary", {}) + print_summary_panel( + total=summary.get("total", 0), + successful=summary.get("successful", 0), + failed=summary.get("failed", 0), + duration=summary.get("duration", 0.0), + ) + else: + print_warning("No test results generated") + + finally: + await orchestrator.cleanup() + async def _run_config_scan(config_path: Path, verbose: bool) -> None: """ @@ -246,8 +247,59 @@ async def _run_config_scan(config_path: Path, verbose: bool) -> None: verbose: Enable verbose output """ print_info(f"Loading configuration from {config_path}") - # TODO: Integrate with core.tester.PromptInjectionTester - print_warning("Config-based scanning not yet implemented") + + # Create orchestrator (URL will be loaded from config) + orchestrator = WorkflowOrchestrator( + target_url="", # Will be overridden by config + verbose=verbose, + ) + + try: + # Run workflow with progress tracking + with create_progress_bar() as progress: + task = progress.add_task( + "Running configuration-based scan...", + total=100, + ) + + # Execute workflow + workflow_results = await orchestrator.run_config_workflow( + config_path=config_path, + progress=progress, + progress_task=task, + ) + + progress.update(task, completed=100) + + # Check for errors + if workflow_results.get("errors"): + for error in workflow_results["errors"]: + print_error(f"Error: {error}") + + if not workflow_results.get("success"): + return + + # Display results + results = workflow_results.get("tests", []) + if results: + console.print() + table = create_results_table(results) + console.print(table) + + # Display summary + console.print() + summary = workflow_results.get("summary", {}) + print_summary_panel( + total=summary.get("total", 0), + successful=summary.get("successful", 0), + failed=summary.get("failed", 0), + duration=summary.get("duration", 0.0), + ) + else: + print_warning("No test results generated") + + finally: + await orchestrator.cleanup() async def _run_manual_scan( diff --git a/tools/prompt_injection_tester/pit/orchestrator/discovery.py b/tools/prompt_injection_tester/pit/orchestrator/discovery.py new file mode 100644 index 0000000..2506405 --- /dev/null +++ b/tools/prompt_injection_tester/pit/orchestrator/discovery.py @@ -0,0 +1,123 @@ +""" +Auto-discovery functionality for LLM endpoints. + +Provides smart detection of models and API capabilities. +""" + +from typing import Optional +from urllib.parse import urlparse + +# Handle optional imports +try: + import aiohttp +except ImportError: + aiohttp = None # type: ignore + + +class AutoDiscovery: + """ + Auto-discovery for LLM endpoints. + + Attempts to detect model identifiers and API capabilities. + """ + + def __init__(self, target_url: str): + """ + Initialize auto-discovery. + + Args: + target_url: Target API endpoint URL + """ + self.target_url = target_url + self.parsed_url = urlparse(target_url) + + async def discover_model(self, hint: Optional[str] = None) -> str: + """ + Discover the model identifier. + + Args: + hint: Optional model hint from user + + Returns: + Discovered or default model identifier + """ + if hint: + return hint + + # Try to discover from endpoint + discovered = await self._probe_for_model() + if discovered: + return discovered + + # Default models based on port/URL + return self._guess_model_from_url() + + async def _probe_for_model(self) -> Optional[str]: + """ + Probe the endpoint for model information. + + Returns: + Model identifier if discovered, None otherwise + """ + if not aiohttp: + return None + + try: + # Try Ollama tags endpoint + if "11434" in self.target_url: + tags_url = f"{self.parsed_url.scheme}://{self.parsed_url.netloc}/api/tags" + async with aiohttp.ClientSession() as session: + async with session.get(tags_url, timeout=aiohttp.ClientTimeout(total=5)) as resp: + if resp.status == 200: + data = await resp.json() + if "models" in data and len(data["models"]) > 0: + return data["models"][0].get("name", "") + except Exception: + pass + + return None + + def _guess_model_from_url(self) -> str: + """ + Guess model based on URL patterns. + + Returns: + Best guess for model identifier + """ + url_lower = self.target_url.lower() + + # Ollama default + if "11434" in url_lower or "ollama" in url_lower: + return "llama3:latest" + + # OpenAI + if "openai" in url_lower or "api.openai.com" in url_lower: + return "gpt-4" + + # Anthropic + if "anthropic" in url_lower or "claude" in url_lower: + return "claude-3-5-sonnet-20241022" + + # LocalAI + if "8080" in url_lower or "localai" in url_lower: + return "gpt-3.5-turbo" + + # Default + return "gpt-3.5-turbo" + + async def discover_capabilities(self) -> dict: + """ + Discover API capabilities. + + Returns: + Dictionary of discovered capabilities + """ + capabilities = { + "streaming": False, + "tool_use": False, + "vision": False, + "embeddings": False, + } + + # This could be expanded with actual probing + return capabilities diff --git a/tools/prompt_injection_tester/pit/orchestrator/workflow.py b/tools/prompt_injection_tester/pit/orchestrator/workflow.py new file mode 100644 index 0000000..279ef53 --- /dev/null +++ b/tools/prompt_injection_tester/pit/orchestrator/workflow.py @@ -0,0 +1,251 @@ +""" +Workflow orchestrator for coordinating scan operations. + +Bridges the CLI layer with the core testing framework. +""" + +import asyncio +import sys +from pathlib import Path +from typing import Optional, List, Dict, Any +from datetime import datetime + +# Handle optional imports +try: + from rich.progress import Progress +except ImportError: + Progress = None # type: ignore + +# Import core framework components +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +from prompt_injection_tester.core.tester import InjectionTester +from prompt_injection_tester.core.models import ( + TargetConfig, + AttackConfig, + TestSuite, + TestResult, +) +from prompt_injection_tester.patterns.registry import registry as pattern_registry + +from pit.orchestrator.discovery import AutoDiscovery + + +class WorkflowOrchestrator: + """ + Orchestrates the complete testing workflow. + + Coordinates between CLI UI and core framework functionality. + """ + + def __init__( + self, + target_url: str, + model: Optional[str] = None, + auth_token: Optional[str] = None, + verbose: bool = False, + ): + """ + Initialize the orchestrator. + + Args: + target_url: Target API endpoint URL + model: Optional model identifier + auth_token: Optional authentication token + verbose: Enable verbose logging + """ + self.target_url = target_url + self.model = model + self.auth_token = auth_token or "" + self.verbose = verbose + self.tester: Optional[InjectionTester] = None + self.discovery = AutoDiscovery(target_url) + + async def run_auto_workflow( + self, + patterns: Optional[List[str]] = None, + progress: Optional[Progress] = None, + progress_task: Optional[Any] = None, + ) -> Dict[str, Any]: + """ + Run automatic workflow with discovery and testing. + + Args: + patterns: Optional list of pattern IDs to test + progress: Optional Rich progress bar + progress_task: Optional progress task ID + + Returns: + Dictionary with workflow results + """ + results = { + "success": False, + "target": self.target_url, + "model": self.model, + "discovery": {}, + "tests": [], + "summary": {}, + "errors": [], + } + + try: + # Phase 1: Discovery + if progress and progress_task is not None: + progress.update(progress_task, description="Discovering endpoint...") + + discovered_model = await self.discovery.discover_model(self.model) + results["model"] = discovered_model + results["discovery"]["model"] = discovered_model + + # Phase 2: Pattern Loading + if not patterns: + patterns = [ + "direct_instruction_override", + "direct_role_authority", + "direct_persona_shift", + ] + + if progress and progress_task is not None: + progress.update( + progress_task, + description=f"Loaded {len(patterns)} patterns", + ) + + # Phase 3: Initialize Tester + target_config = TargetConfig( + name="CLI Target", + base_url=self.target_url, + api_type="openai", + model=discovered_model, + auth_token=self.auth_token, + timeout=30, + rate_limit=1.0, + ) + + attack_config = AttackConfig( + patterns=patterns, + max_concurrent=3, + timeout_per_test=30, + rate_limit=1.0, + stop_on_success=False, + ) + + self.tester = InjectionTester( + target_config=target_config, + config=attack_config, + ) + + # Grant authorization + self.tester.authorize(["all"]) + + # Phase 4: Run Tests + async with self.tester: + test_suite = await self.tester.run_tests(patterns=patterns) + + # Convert results + test_results = [] + for test_result in test_suite.results: + test_results.append({ + "pattern": test_result.pattern_id, + "success": test_result.status.value == "success", + "confidence": test_result.confidence, + "details": test_result.evidence[:100] if test_result.evidence else "", + "response": test_result.response[:200] if test_result.response else "", + }) + + results["tests"] = test_results + results["summary"] = { + "total": len(test_results), + "successful": sum(1 for r in test_results if r["success"]), + "failed": sum(1 for r in test_results if not r["success"]), + "duration": ( + (test_suite.completed_at - test_suite.started_at).total_seconds() + if test_suite.completed_at and test_suite.started_at + else 0.0 + ), + } + results["success"] = True + + except Exception as e: + results["errors"].append(str(e)) + if self.verbose: + import traceback + results["errors"].append(traceback.format_exc()) + + return results + + async def run_config_workflow( + self, + config_path: Path, + progress: Optional[Progress] = None, + progress_task: Optional[Any] = None, + ) -> Dict[str, Any]: + """ + Run workflow using configuration file. + + Args: + config_path: Path to YAML configuration + progress: Optional Rich progress bar + progress_task: Optional progress task ID + + Returns: + Dictionary with workflow results + """ + results = { + "success": False, + "config": str(config_path), + "tests": [], + "summary": {}, + "errors": [], + } + + try: + if progress and progress_task is not None: + progress.update(progress_task, description="Loading configuration...") + + # Load tester from config + self.tester = InjectionTester.from_config_file(config_path) + self.tester.authorize(["all"]) + + if progress and progress_task is not None: + progress.update(progress_task, description="Running tests...") + + # Run tests + async with self.tester: + test_suite = await self.tester.run_tests() + + # Convert results + test_results = [] + for test_result in test_suite.results: + test_results.append({ + "pattern": test_result.pattern_id, + "success": test_result.status.value == "success", + "confidence": test_result.confidence, + "details": test_result.evidence[:100] if test_result.evidence else "", + }) + + results["tests"] = test_results + results["summary"] = { + "total": len(test_results), + "successful": sum(1 for r in test_results if r["success"]), + "failed": sum(1 for r in test_results if not r["success"]), + "duration": ( + (test_suite.completed_at - test_suite.started_at).total_seconds() + if test_suite.completed_at and test_suite.started_at + else 0.0 + ), + } + results["success"] = True + + except Exception as e: + results["errors"].append(str(e)) + if self.verbose: + import traceback + results["errors"].append(traceback.format_exc()) + + return results + + async def cleanup(self): + """Cleanup resources.""" + if self.tester: + await self.tester.close()