feat: implement orchestrator module to manage LLM endpoint discovery and testing workflows.

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(

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

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()