zen-ai-pentest 2.2.0__py3-none-any.whl → 2.3.0__py3-none-any.whl

benchmarks/scan_performance.py

@@ -0,0 +1,272 @@
+"""
+Scan Performance Benchmarks
+
+Measures scan throughput, target processing speed, and tool execution times.
+"""
+
+import asyncio
+import time
+from typing import List, Dict, Any, Optional, Callable
+from dataclasses import dataclass, field
+
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from modules.benchmark import (
+    BenchmarkRunner, BenchmarkResult, BenchmarkCategory,
+    ThroughputMetrics, measure_scan_throughput
+)
+
+
+@dataclass
+class ScanBenchmarkConfig:
+    """Configuration for scan benchmarks."""
+    target_count: int = 10
+    target_type: str = "host"  # host, web, network
+    ports: str = "top-100"
+    scan_depth: str = "quick"  # quick, standard, deep
+    concurrent_scans: int = 3
+    warmup_targets: int = 2
+
+
+class ScanPerformanceBenchmark:
+    """Benchmark suite for scan performance measurement."""
+    
+    def __init__(self, output_dir: str = "benchmark_results"):
+        self.runner = BenchmarkRunner(output_dir=output_dir)
+        self.config = ScanBenchmarkConfig()
+    
+    def _generate_test_targets(self, count: int, target_type: str) -> List[str]:
+        """Generate test targets for benchmarking."""
+        if target_type == "host":
+            # Generate local network targets for testing
+            return [f"192.168.1.{i}" for i in range(1, min(count + 1, 255))]
+        elif target_type == "web":
+            # Generate test web targets
+            return [f"test-target-{i}.local" for i in range(count)]
+        elif target_type == "network":
+            # Generate network ranges
+            ranges = []
+            for i in range(min(count, 10)):
+                ranges.append(f"10.0.{i}.0/24")
+            return ranges
+        return []
+    
+    async def benchmark_nmap_speed(self, config: Optional[ScanBenchmarkConfig] = None) -> BenchmarkResult:
+        """
+        Benchmark Nmap scan speed.
+        
+        Returns:
+            BenchmarkResult with scan timing metrics
+        """
+        cfg = config or self.config
+        targets = self._generate_test_targets(cfg.target_count, cfg.target_type)
+        
+        # Mock scan function for benchmarking structure
+        async def mock_nmap_scan(target: str) -> Dict[str, Any]:
+            # Simulate scan delay
+            await asyncio.sleep(0.5)
+            return {"target": target, "status": "completed", "ports_found": 5}
+        
+        result, _ = await measure_scan_throughput(
+            scan_func=mock_nmap_scan,
+            targets=targets,
+        )
+        
+        result.name = "nmap_scan_speed"
+        result.description = f"Nmap scan speed for {len(targets)} targets"
+        
+        return result
+    
+    async def benchmark_web_scan_speed(self, config: Optional[ScanBenchmarkConfig] = None) -> BenchmarkResult:
+        """
+        Benchmark web vulnerability scan speed.
+        
+        Returns:
+            BenchmarkResult with web scan timing metrics
+        """
+        cfg = config or ScanBenchmarkConfig(target_count=5, target_type="web")
+        targets = self._generate_test_targets(cfg.target_count, cfg.target_type)
+        
+        async def mock_web_scan(target: str) -> Dict[str, Any]:
+            # Simulate web scan with crawling
+            await asyncio.sleep(1.2)
+            return {
+                "target": target,
+                "pages_scanned": 15,
+                "vulnerabilities_found": 3
+            }
+        
+        result, _ = await measure_scan_throughput(
+            scan_func=mock_web_scan,
+            targets=targets,
+        )
+        
+        result.name = "web_scan_speed"
+        result.description = f"Web vulnerability scan speed for {len(targets)} targets"
+        
+        return result
+    
+    async def benchmark_scan_throughput(self, config: Optional[ScanBenchmarkConfig] = None) -> BenchmarkResult:
+        """
+        Benchmark overall scan throughput (targets per minute).
+        
+        Returns:
+            BenchmarkResult with throughput metrics
+        """
+        cfg = config or self.config
+        targets = self._generate_test_targets(cfg.target_count, cfg.target_type)
+        
+        async def mock_scan(target: str) -> Dict[str, Any]:
+            # Simulate varying scan times
+            delay = 0.3 + (hash(target) % 10) / 10  # 0.3-1.3s
+            await asyncio.sleep(delay)
+            return {"target": target, "completed": True}
+        
+        result, _ = await measure_scan_throughput(
+            scan_func=mock_scan,
+            targets=targets,
+        )
+        
+        result.name = "scan_throughput"
+        result.description = f"Overall scan throughput: {len(targets)} targets"
+        
+        return result
+    
+    async def benchmark_concurrent_scans(self, config: Optional[ScanBenchmarkConfig] = None) -> BenchmarkResult:
+        """
+        Benchmark performance with concurrent scans.
+        
+        Returns:
+            BenchmarkResult with concurrent scan metrics
+        """
+        cfg = config or self.config
+        targets = self._generate_test_targets(cfg.target_count * 2, cfg.target_type)
+        
+        async def concurrent_benchmark():
+            semaphore = asyncio.Semaphore(cfg.concurrent_scans)
+            
+            async def scan_with_limit(target: str):
+                async with semaphore:
+                    await asyncio.sleep(0.4)
+                    return {"target": target}
+            
+            await asyncio.gather(*[scan_with_limit(t) for t in targets])
+        
+        result = await self.runner.run_benchmark(
+            name="concurrent_scan_performance",
+            category=BenchmarkCategory.SCAN,
+            description=f"Concurrent scan performance ({cfg.concurrent_scans} parallel)",
+            benchmark_func=concurrent_benchmark,
+            iterations=1,
+            monitor_resources=True
+        )
+        
+        # Calculate effective throughput
+        result.throughput = ThroughputMetrics(
+            operations=len(targets),
+            duration_seconds=result.timing.duration_seconds
+        )
+        result.throughput.calculate()
+        
+        result.custom_metrics["concurrent_scans"] = cfg.concurrent_scans
+        result.custom_metrics["total_targets"] = len(targets)
+        
+        return result
+    
+    async def run_all(self) -> List[BenchmarkResult]:
+        """Run all scan performance benchmarks."""
+        results = []
+        
+        print("Running scan performance benchmarks...")
+        
+        # Run each benchmark
+        benchmarks = [
+            ("Nmap Speed", self.benchmark_nmap_speed),
+            ("Web Scan Speed", self.benchmark_web_scan_speed),
+            ("Scan Throughput", self.benchmark_scan_throughput),
+            ("Concurrent Scans", self.benchmark_concurrent_scans),
+        ]
+        
+        for name, benchmark_func in benchmarks:
+            print(f"  Running: {name}...")
+            try:
+                result = await benchmark_func()
+                results.append(result)
+                self.runner.save_result(result)
+                print(f"    ✓ {name}: {result.throughput.ops_per_minute:.1f} targets/min")
+            except Exception as e:
+                print(f"    ✗ {name} failed: {e}")
+        
+        # Save combined results
+        self.runner.save_all_results("scan_benchmarks.json")
+        
+        return results
+    
+    def get_summary(self) -> Dict[str, Any]:
+        """Get summary of scan benchmark results."""
+        return self.runner.get_summary()
+
+
+# Convenience function
+async def measure_scan_speed(
+    scan_func: Callable,
+    targets: List[str],
+    output_dir: str = "benchmark_results"
+) -> BenchmarkResult:
+    """
+    Quick function to measure scan speed.
+    
+    Args:
+        scan_func: Async function to perform scan
+        targets: List of targets
+        output_dir: Directory for results
+    
+    Returns:
+        BenchmarkResult with scan metrics
+    """
+    result, _ = await measure_scan_throughput(
+        scan_func=scan_func,
+        targets=targets
+    )
+    
+    # Save result
+    runner = BenchmarkRunner(output_dir=output_dir)
+    runner.save_result(result)
+    
+    return result
+
+
+# CLI interface
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="Scan Performance Benchmarks")
+    parser.add_argument("--output", default="benchmark_results", help="Output directory")
+    parser.add_argument("--targets", type=int, default=10, help="Number of test targets")
+    parser.add_argument("--type", default="host", choices=["host", "web", "network"],
+                       help="Target type")
+    
+    args = parser.parse_args()
+    
+    async def main():
+        config = ScanBenchmarkConfig(
+            target_count=args.targets,
+            target_type=args.type
+        )
+        
+        benchmark = ScanPerformanceBenchmark(output_dir=args.output)
+        results = await benchmark.run_all()
+        
+        print("\n" + "="*60)
+        print("SCAN PERFORMANCE BENCHMARK RESULTS")
+        print("="*60)
+        
+        for result in results:
+            print(f"\n{result.name}:")
+            print(f"  Duration: {result.timing.duration_ms:.2f}ms")
+            print(f"  Throughput: {result.throughput.ops_per_minute:.2f} ops/min")
+            print(f"  Peak Memory: {result.memory.peak_mb:.2f} MB")
+    
+    asyncio.run(main())

modules/agent_coordinator.py

@@ -0,0 +1,255 @@
+"""Multi-Agent Coordinator with Deadlock Prevention
+
+Prevents deadlocks in multi-agent systems using resource ordering and timeouts.
+Addresses Issue #13
+"""
+import asyncio
+import time
+from typing import Dict, List, Set, Optional, Any
+from dataclasses import dataclass, field
+from enum import Enum
+from contextlib import asynccontextmanager
+import logging
+
+
+class AgentStatus(Enum):
+    IDLE = "idle"
+    RUNNING = "running"
+    WAITING = "waiting"
+    COMPLETED = "completed"
+    ERROR = "error"
+
+
+class ResourceType(Enum):
+    SCANNER = "scanner"
+    DATABASE = "database"
+    API_RATE_LIMIT = "api_rate_limit"
+    FILE_SYSTEM = "file_system"
+    NETWORK = "network"
+
+
+@dataclass
+class Agent:
+    """Agent representation"""
+    id: str
+    name: str
+    status: AgentStatus = AgentStatus.IDLE
+    acquired_resources: Set[ResourceType] = field(default_factory=set)
+    waiting_for: Optional[ResourceType] = None
+    start_time: float = field(default_factory=time.time)
+    timeout: float = 300.0  # 5 minutes default
+
+
+@dataclass
+class Resource:
+    """Resource representation"""
+    type: ResourceType
+    max_concurrent: int = 1
+    current_users: Set[str] = field(default_factory=set)
+    wait_queue: List[str] = field(default_factory=list)
+
+
+class AgentCoordinator:
+    """Coordinates multiple agents with deadlock prevention"""
+
+    name = "agent_coordinator"
+    version = "1.0.0"
+
+    def __init__(self):
+        self.agents: Dict[str, Agent] = {}
+        self.resources: Dict[ResourceType, Resource] = {}
+        self._lock = asyncio.Lock()
+        self._init_resources()
+        self.deadlock_check_interval = 10.0
+
+    def _init_resources(self):
+        """Initialize system resources"""
+        self.resources = {
+            ResourceType.SCANNER: Resource(ResourceType.SCANNER, max_concurrent=3),
+            ResourceType.DATABASE: Resource(ResourceType.DATABASE, max_concurrent=5),
+            ResourceType.API_RATE_LIMIT: Resource(ResourceType.API_RATE_LIMIT, max_concurrent=2),
+            ResourceType.FILE_SYSTEM: Resource(ResourceType.FILE_SYSTEM, max_concurrent=10),
+            ResourceType.NETWORK: Resource(ResourceType.NETWORK, max_concurrent=5),
+        }
+
+    async def register_agent(self, agent_id: str, name: str, timeout: float = 300.0) -> Agent:
+        """Register a new agent"""
+        async with self._lock:
+            if agent_id in self.agents:
+                raise ValueError(f"Agent {agent_id} already registered")
+            
+            agent = Agent(id=agent_id, name=name, timeout=timeout)
+            self.agents[agent_id] = agent
+            return agent
+
+    async def unregister_agent(self, agent_id: str):
+        """Unregister an agent and release its resources"""
+        async with self._lock:
+            if agent_id not in self.agents:
+                return
+            
+            agent = self.agents[agent_id]
+            
+            # Release all acquired resources
+            for resource_type in list(agent.acquired_resources):
+                await self._release_resource(agent_id, resource_type)
+            
+            del self.agents[agent_id]
+
+    async def _acquire_resource(
+        self,
+        agent_id: str,
+        resource_type: ResourceType,
+        timeout: float = 30.0
+    ) -> bool:
+        """Acquire a resource with timeout"""
+        start_time = time.time()
+        
+        while time.time() - start_time < timeout:
+            async with self._lock:
+                if agent_id not in self.agents:
+                    return False
+                
+                resource = self.resources[resource_type]
+                agent = self.agents[agent_id]
+                
+                # Check if resource available
+                if len(resource.current_users) < resource.max_concurrent:
+                    resource.current_users.add(agent_id)
+                    agent.acquired_resources.add(resource_type)
+                    agent.waiting_for = None
+                    return True
+                
+                # Add to wait queue
+                if agent_id not in resource.wait_queue:
+                    resource.wait_queue.append(agent_id)
+                    agent.waiting_for = resource_type
+                    agent.status = AgentStatus.WAITING
+            
+            # Wait and retry
+            await asyncio.sleep(0.5)
+        
+        return False
+
+    async def _release_resource(self, agent_id: str, resource_type: ResourceType):
+        """Release an acquired resource"""
+        resource = self.resources[resource_type]
+        
+        if agent_id in resource.current_users:
+            resource.current_users.remove(agent_id)
+        
+        if agent_id in resource.wait_queue:
+            resource.wait_queue.remove(agent_id)
+        
+        if agent_id in self.agents:
+            agent = self.agents[agent_id]
+            agent.acquired_resources.discard(resource_type)
+
+    @asynccontextmanager
+    async def acquire_resources(
+        self,
+        agent_id: str,
+        resources: List[ResourceType],
+        timeout: float = 30.0
+    ):
+        """
+        Context manager for acquiring multiple resources
+        Uses resource ordering to prevent deadlocks
+        """
+        # Sort resources to ensure consistent ordering (prevents circular wait)
+        sorted_resources = sorted(resources, key=lambda r: r.value)
+        
+        acquired = []
+        try:
+            # Acquire all resources in order
+            for resource_type in sorted_resources:
+                success = await self._acquire_resource(agent_id, resource_type, timeout)
+                if not success:
+                    raise TimeoutError(f"Could not acquire {resource_type.value}")
+                acquired.append(resource_type)
+            
+            yield acquired
+        finally:
+            # Release in reverse order
+            for resource_type in reversed(acquired):
+                await self._release_resource(agent_id, resource_type)
+
+    async def check_deadlocks(self) -> List[Dict]:
+        """Detect potential deadlocks in the system"""
+        async with self._lock:
+            deadlocks = []
+            
+            for agent_id, agent in self.agents.items():
+                # Check for timeout
+                if agent.status == AgentStatus.WAITING:
+                    wait_time = time.time() - agent.start_time
+                    if wait_time > agent.timeout:
+                        deadlocks.append({
+                            'agent_id': agent_id,
+                            'type': 'timeout',
+                            'wait_time': wait_time,
+                            'waiting_for': agent.waiting_for.value if agent.waiting_for else None
+                        })
+                
+                # Check for circular wait (simplified)
+                if agent.waiting_for:
+                    resource = self.resources[agent.waiting_for]
+                    for other_id in resource.current_users:
+                        if other_id in self.agents:
+                            other = self.agents[other_id]
+                            if other.waiting_for in agent.acquired_resources:
+                                deadlocks.append({
+                                    'agent_id': agent_id,
+                                    'type': 'circular_wait',
+                                    'involved': [agent_id, other_id],
+                                    'resources': [agent.waiting_for.value, other.waiting_for.value if other.waiting_for else None]
+                                })
+            
+            return deadlocks
+
+    async def resolve_deadlock(self, agent_id: str):
+        """Forcefully resolve a deadlock by killing an agent"""
+        async with self._lock:
+            if agent_id not in self.agents:
+                return
+            
+            agent = self.agents[agent_id]
+            
+            # Release all resources
+            for resource_type in list(agent.acquired_resources):
+                await self._release_resource(agent_id, resource_type)
+            
+            agent.status = AgentStatus.ERROR
+            logging.warning(f"Deadlock resolved: Agent {agent_id} terminated")
+
+    def get_status(self) -> Dict:
+        """Get coordinator status"""
+        return {
+            'agents': {
+                aid: {
+                    'name': a.name,
+                    'status': a.status.value,
+                    'resources': [r.value for r in a.acquired_resources],
+                    'waiting_for': a.waiting_for.value if a.waiting_for else None
+                }
+                for aid, a in self.agents.items()
+            },
+            'resources': {
+                rt.value: {
+                    'max': r.max_concurrent,
+                    'current': len(r.current_users),
+                    'queue': len(r.wait_queue)
+                }
+                for rt, r in self.resources.items()
+            }
+        }
+
+    def get_info(self) -> Dict:
+        """Get module info"""
+        return {
+            'name': self.name,
+            'version': self.version,
+            'description': 'Multi-agent coordinator with deadlock prevention',
+            'resources': [r.value for r in ResourceType],
+            'deadlock_prevention': ['resource_ordering', 'timeout_detection', 'circular_wait_detection']
+        }