""" Advanced Multi-Agent Scenario Orchestration for Cyber-LLM Handles complex red team exercises and coordinated agent operations """ import asyncio import json import yaml from typing import Dict, List, Any, Optional, Tuple, Union from datetime import datetime, timedelta from pathlib import Path from dataclasses import dataclass, asdict from enum import Enum import uuid import networkx as nx from concurrent.futures import ThreadPoolExecutor, as_completed from ..agents.orchestrator_agent import OrchestratorAgent from ..agents.recon_agent import ReconnaissanceAgent from ..agents.c2_agent import CommandControlAgent from ..agents.post_exploit_agent import PostExploitAgent from ..agents.safety_agent import SafetyAgent from ..agents.explainability_agent import ExplainabilityAgent from ..utils.logging_system import CyberLLMLogger, CyberLLMError, ErrorCategory, retry_with_backoff class ScenarioType(Enum): """Types of security scenarios""" RED_TEAM_EXERCISE = "red_team_exercise" PENETRATION_TEST = "penetration_test" THREAT_HUNTING = "threat_hunting" INCIDENT_RESPONSE = "incident_response" VULNERABILITY_ASSESSMENT = "vulnerability_assessment" SECURITY_AUDIT = "security_audit" class AgentRole(Enum): """Agent roles in scenarios""" LEADER = "leader" SPECIALIST = "specialist" SUPPORT = "support" OBSERVER = "observer" @dataclass class ScenarioStep: """Individual step in a security scenario""" id: str name: str description: str agent_type: str dependencies: List[str] parameters: Dict[str, Any] timeout: int = 300 # 5 minutes default retry_count: int = 3 critical: bool = False parallel_group: Optional[str] = None @dataclass class ScenarioResult: """Result from executing a scenario step""" step_id: str success: bool output: Any error: Optional[str] execution_time: float timestamp: datetime @dataclass class RedTeamScenario: """Complete red team exercise scenario""" id: str name: str description: str scenario_type: ScenarioType target_environment: Dict[str, Any] steps: List[ScenarioStep] success_criteria: Dict[str, Any] safety_constraints: List[str] estimated_duration: int # minutes difficulty_level: str # beginner, intermediate, advanced, expert class MultiAgentOrchestrator: """Advanced orchestrator for complex multi-agent scenarios""" def __init__(self, logger: Optional[CyberLLMLogger] = None, max_concurrent_agents: int = 5): self.logger = logger or CyberLLMLogger(name="multi_agent_orchestrator") self.max_concurrent_agents = max_concurrent_agents # Initialize agents self.agents = { 'orchestrator': OrchestratorAgent(logger=self.logger), 'recon': ReconnaissanceAgent(logger=self.logger), 'c2': CommandControlAgent(logger=self.logger), 'post_exploit': PostExploitAgent(logger=self.logger), 'safety': SafetyAgent(logger=self.logger), 'explainability': ExplainabilityAgent(logger=self.logger) } # Execution state self.active_scenarios = {} self.scenario_results = {} self.agent_status = {name: "idle" for name in self.agents.keys()} # Scenario templates self.scenario_templates = self._load_scenario_templates() def _load_scenario_templates(self) -> Dict[str, RedTeamScenario]: """Load predefined scenario templates""" templates = {} # Advanced Persistent Threat (APT) Simulation apt_scenario = RedTeamScenario( id="apt_simulation_001", name="Advanced Persistent Threat Simulation", description="Multi-stage APT attack simulation with stealth focus", scenario_type=ScenarioType.RED_TEAM_EXERCISE, target_environment={ "network_range": "10.0.0.0/24", "domain": "target.local", "critical_assets": ["domain_controller", "file_server", "database"] }, steps=[ ScenarioStep( id="recon_phase", name="Reconnaissance", description="Gather intelligence on target environment", agent_type="recon", dependencies=[], parameters={ "target": "target.local", "scope": "external", "stealth_level": "high" }, timeout=600, critical=True ), ScenarioStep( id="initial_access", name="Initial Access", description="Gain initial foothold in target network", agent_type="c2", dependencies=["recon_phase"], parameters={ "attack_vector": "spear_phishing", "payload_type": "custom_backdoor" }, timeout=900, critical=True ), ScenarioStep( id="privilege_escalation", name="Privilege Escalation", description="Escalate privileges on compromised system", agent_type="post_exploit", dependencies=["initial_access"], parameters={ "technique": "token_manipulation", "target_privilege": "SYSTEM" }, timeout=300 ), ScenarioStep( id="lateral_movement", name="Lateral Movement", description="Move laterally through the network", agent_type="post_exploit", dependencies=["privilege_escalation"], parameters={ "technique": "pass_the_hash", "targets": ["file_server", "database"] }, timeout=600, parallel_group="expansion" ), ScenarioStep( id="data_exfiltration", name="Data Exfiltration", description="Exfiltrate sensitive data", agent_type="post_exploit", dependencies=["lateral_movement"], parameters={ "data_types": ["financial", "pii", "intellectual_property"], "exfil_method": "dns_tunneling" }, timeout=450, parallel_group="expansion" ), ScenarioStep( id="persistence", name="Establish Persistence", description="Maintain long-term access", agent_type="c2", dependencies=["privilege_escalation"], parameters={ "persistence_type": "service_persistence", "stealth_level": "high" }, timeout=300 ) ], success_criteria={ "min_steps_completed": 4, "critical_steps_required": ["recon_phase", "initial_access"], "stealth_score_threshold": 0.7, "detection_threshold": 0.3 }, safety_constraints=[ "no_actual_data_modification", "no_service_disruption", "maintain_system_stability", "document_all_actions" ], estimated_duration=180, # 3 hours difficulty_level="advanced" ) templates["apt_simulation"] = apt_scenario # Vulnerability Assessment Scenario vuln_assessment = RedTeamScenario( id="vuln_assessment_001", name="Comprehensive Vulnerability Assessment", description="Multi-vector vulnerability discovery and assessment", scenario_type=ScenarioType.VULNERABILITY_ASSESSMENT, target_environment={ "targets": ["web_app", "network_services", "host_systems"], "assessment_scope": "comprehensive" }, steps=[ ScenarioStep( id="network_discovery", name="Network Discovery", description="Discover network topology and services", agent_type="recon", dependencies=[], parameters={ "scan_type": "comprehensive", "port_range": "1-65535" }, timeout=1800, parallel_group="discovery" ), ScenarioStep( id="service_enumeration", name="Service Enumeration", description="Enumerate discovered services", agent_type="recon", dependencies=[], parameters={ "service_types": ["web", "database", "file_sharing"], "deep_scan": True }, timeout=1200, parallel_group="discovery" ), ScenarioStep( id="vulnerability_scanning", name="Vulnerability Scanning", description="Scan for known vulnerabilities", agent_type="recon", dependencies=["network_discovery", "service_enumeration"], parameters={ "scanner_types": ["nessus", "openvas", "custom"], "authenticated": False }, timeout=2400, critical=True ), ScenarioStep( id="web_app_testing", name="Web Application Testing", description="Test web applications for vulnerabilities", agent_type="recon", dependencies=["service_enumeration"], parameters={ "test_types": ["owasp_top10", "custom_checks"], "authentication_bypass": True }, timeout=1800, parallel_group="testing" ), ScenarioStep( id="exploit_validation", name="Exploit Validation", description="Validate critical vulnerabilities", agent_type="c2", dependencies=["vulnerability_scanning"], parameters={ "exploit_types": ["proof_of_concept"], "severity_threshold": "high" }, timeout=900 ) ], success_criteria={ "vulnerability_discovery_rate": 0.8, "false_positive_rate": 0.1, "coverage_percentage": 0.9 }, safety_constraints=[ "read_only_operations", "no_system_modification", "minimal_service_impact" ], estimated_duration=360, # 6 hours difficulty_level="intermediate" ) templates["vuln_assessment"] = vuln_assessment return templates async def execute_scenario(self, scenario: RedTeamScenario, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: """Execute a complete multi-agent scenario""" scenario_id = f"{scenario.id}_{uuid.uuid4().hex[:8]}" start_time = datetime.now() self.logger.info(f"Starting scenario execution", scenario_id=scenario_id, scenario_name=scenario.name, scenario_type=scenario.scenario_type.value, estimated_duration=scenario.estimated_duration) # Initialize scenario state self.active_scenarios[scenario_id] = { 'scenario': scenario, 'context': context or {}, 'start_time': start_time, 'status': 'running', 'completed_steps': [], 'failed_steps': [], 'step_results': {} } try: # Safety check safety_approval = await self._safety_check(scenario) if not safety_approval['approved']: raise CyberLLMError( f"Scenario failed safety check: {safety_approval['reason']}", ErrorCategory.SAFETY ) # Build execution graph execution_graph = self._build_execution_graph(scenario.steps) # Execute scenario steps results = await self._execute_scenario_graph( scenario_id, execution_graph, scenario.success_criteria ) # Evaluate results evaluation = self._evaluate_scenario_results(scenario, results) # Generate report report = await self._generate_scenario_report( scenario_id, scenario, results, evaluation ) execution_time = (datetime.now() - start_time).total_seconds() self.active_scenarios[scenario_id]['status'] = 'completed' self.scenario_results[scenario_id] = { 'scenario': scenario, 'results': results, 'evaluation': evaluation, 'report': report, 'execution_time': execution_time, 'completed_at': datetime.now() } self.logger.info(f"Scenario execution completed", scenario_id=scenario_id, success=evaluation['overall_success'], execution_time=execution_time, steps_completed=len(results['completed']), steps_failed=len(results['failed'])) return self.scenario_results[scenario_id] except Exception as e: self.active_scenarios[scenario_id]['status'] = 'failed' self.logger.error(f"Scenario execution failed", scenario_id=scenario_id, error=str(e)) raise finally: # Cleanup if scenario_id in self.active_scenarios: del self.active_scenarios[scenario_id] async def _safety_check(self, scenario: RedTeamScenario) -> Dict[str, Any]: """Perform safety check on scenario""" safety_agent = self.agents['safety'] # Check scenario against safety constraints check_request = { 'scenario_type': scenario.scenario_type.value, 'target_environment': scenario.target_environment, 'steps': [asdict(step) for step in scenario.steps], 'safety_constraints': scenario.safety_constraints } try: safety_result = await safety_agent.evaluate_scenario_safety(check_request) return { 'approved': safety_result.get('approved', False), 'reason': safety_result.get('reason', ''), 'risk_level': safety_result.get('risk_level', 'unknown'), 'recommendations': safety_result.get('recommendations', []) } except Exception as e: self.logger.error("Safety check failed", error=str(e)) return { 'approved': False, 'reason': f"Safety check error: {str(e)}", 'risk_level': 'critical' } def _build_execution_graph(self, steps: List[ScenarioStep]) -> nx.DiGraph: """Build directed graph for scenario execution""" graph = nx.DiGraph() # Add nodes for step in steps: graph.add_node(step.id, step=step) # Add dependency edges for step in steps: for dependency in step.dependencies: if dependency in [s.id for s in steps]: graph.add_edge(dependency, step.id) # Verify graph is acyclic if not nx.is_directed_acyclic_graph(graph): raise CyberLLMError( "Scenario contains circular dependencies", ErrorCategory.VALIDATION ) return graph async def _execute_scenario_graph(self, scenario_id: str, graph: nx.DiGraph, success_criteria: Dict[str, Any]) -> Dict[str, Any]: """Execute scenario steps according to dependency graph""" completed_steps = set() failed_steps = set() step_results = {} parallel_groups = {} # Group steps by parallel execution groups for node_id in graph.nodes(): step = graph.nodes[node_id]['step'] if step.parallel_group: if step.parallel_group not in parallel_groups: parallel_groups[step.parallel_group] = [] parallel_groups[step.parallel_group].append(step) # Execute steps in topological order execution_order = list(nx.topological_sort(graph)) for step_id in execution_order: step = graph.nodes[step_id]['step'] # Check if dependencies are satisfied dependencies_met = all( dep in completed_steps for dep in step.dependencies ) if not dependencies_met: failed_steps.add(step_id) self.logger.warning(f"Step dependencies not met: {step_id}") continue # Execute step try: if step.parallel_group and step.parallel_group in parallel_groups: # Execute parallel group group_steps = parallel_groups[step.parallel_group] group_results = await self._execute_parallel_steps(group_steps) for group_step, result in group_results.items(): step_results[group_step.id] = result if result.success: completed_steps.add(group_step.id) else: failed_steps.add(group_step.id) # Remove processed group del parallel_groups[step.parallel_group] else: # Execute single step result = await self._execute_single_step(step) step_results[step_id] = result if result.success: completed_steps.add(step_id) else: failed_steps.add(step_id) # Check if critical step failed if step.critical: self.logger.error(f"Critical step failed: {step_id}") break except Exception as e: self.logger.error(f"Step execution error: {step_id}", error=str(e)) failed_steps.add(step_id) if step.critical: break return { 'completed': completed_steps, 'failed': failed_steps, 'results': step_results, 'success_rate': len(completed_steps) / len(graph.nodes()) if graph.nodes() else 0 } async def _execute_parallel_steps(self, steps: List[ScenarioStep]) -> Dict[ScenarioStep, ScenarioResult]: """Execute multiple steps in parallel""" tasks = [] for step in steps: task = asyncio.create_task(self._execute_single_step(step)) tasks.append((step, task)) results = {} completed_tasks = await asyncio.gather(*[task for _, task in tasks], return_exceptions=True) for (step, _), result in zip(tasks, completed_tasks): if isinstance(result, Exception): results[step] = ScenarioResult( step_id=step.id, success=False, output=None, error=str(result), execution_time=0, timestamp=datetime.now() ) else: results[step] = result return results @retry_with_backoff(max_retries=3) async def _execute_single_step(self, step: ScenarioStep) -> ScenarioResult: """Execute a single scenario step""" start_time = datetime.now() self.logger.info(f"Executing step: {step.name}", step_id=step.id, agent_type=step.agent_type) try: # Get appropriate agent agent = self.agents.get(step.agent_type) if not agent: raise CyberLLMError( f"Unknown agent type: {step.agent_type}", ErrorCategory.VALIDATION ) # Update agent status self.agent_status[step.agent_type] = "busy" # Execute step with timeout result = await asyncio.wait_for( agent.execute_task(step.parameters), timeout=step.timeout ) execution_time = (datetime.now() - start_time).total_seconds() step_result = ScenarioResult( step_id=step.id, success=True, output=result, error=None, execution_time=execution_time, timestamp=datetime.now() ) self.logger.info(f"Step completed successfully: {step.name}", step_id=step.id, execution_time=execution_time) return step_result except asyncio.TimeoutError: error_msg = f"Step timed out after {step.timeout} seconds" self.logger.error(f"Step timeout: {step.name}", step_id=step.id) return ScenarioResult( step_id=step.id, success=False, output=None, error=error_msg, execution_time=step.timeout, timestamp=datetime.now() ) except Exception as e: execution_time = (datetime.now() - start_time).total_seconds() error_msg = str(e) self.logger.error(f"Step execution failed: {step.name}", step_id=step.id, error=error_msg) return ScenarioResult( step_id=step.id, success=False, output=None, error=error_msg, execution_time=execution_time, timestamp=datetime.now() ) finally: # Reset agent status self.agent_status[step.agent_type] = "idle" def _evaluate_scenario_results(self, scenario: RedTeamScenario, results: Dict[str, Any]) -> Dict[str, Any]: """Evaluate scenario results against success criteria""" success_criteria = scenario.success_criteria evaluation = { 'overall_success': False, 'criteria_met': {}, 'score': 0.0, 'recommendations': [] } # Check minimum steps completed min_steps = success_criteria.get('min_steps_completed', 0) steps_completed = len(results['completed']) evaluation['criteria_met']['min_steps_completed'] = steps_completed >= min_steps # Check critical steps critical_steps = success_criteria.get('critical_steps_required', []) critical_met = all(step in results['completed'] for step in critical_steps) evaluation['criteria_met']['critical_steps_completed'] = critical_met # Calculate success score total_steps = len(results['completed']) + len(results['failed']) if total_steps > 0: success_rate = len(results['completed']) / total_steps evaluation['score'] = success_rate # Overall success determination evaluation['overall_success'] = ( evaluation['criteria_met'].get('min_steps_completed', False) and evaluation['criteria_met'].get('critical_steps_completed', False) and evaluation['score'] >= 0.7 # 70% success threshold ) # Generate recommendations if not evaluation['overall_success']: if not critical_met: evaluation['recommendations'].append("Complete all critical steps") if evaluation['score'] < 0.7: evaluation['recommendations'].append("Improve step success rate") return evaluation async def _generate_scenario_report(self, scenario_id: str, scenario: RedTeamScenario, results: Dict[str, Any], evaluation: Dict[str, Any]) -> Dict[str, Any]: """Generate comprehensive scenario execution report""" # Get explainability analysis explainability_agent = self.agents['explainability'] analysis_request = { 'scenario': asdict(scenario), 'results': results, 'evaluation': evaluation } try: explanation = await explainability_agent.analyze_scenario_execution(analysis_request) except Exception as e: self.logger.warning("Failed to generate explanation", error=str(e)) explanation = {"analysis": "Analysis unavailable", "insights": []} report = { 'scenario_id': scenario_id, 'scenario_summary': { 'name': scenario.name, 'type': scenario.scenario_type.value, 'difficulty': scenario.difficulty_level, 'estimated_duration': scenario.estimated_duration }, 'execution_summary': { 'total_steps': len(scenario.steps), 'completed_steps': len(results['completed']), 'failed_steps': len(results['failed']), 'success_rate': results['success_rate'], 'overall_success': evaluation['overall_success'] }, 'detailed_results': results['results'], 'evaluation': evaluation, 'explanation': explanation, 'generated_at': datetime.now().isoformat() } return report def get_scenario_template(self, template_name: str) -> Optional[RedTeamScenario]: """Get a scenario template by name""" return self.scenario_templates.get(template_name) def list_scenario_templates(self) -> List[str]: """List available scenario templates""" return list(self.scenario_templates.keys()) def get_active_scenarios(self) -> Dict[str, Any]: """Get currently active scenarios""" return self.active_scenarios.copy() def get_agent_status(self) -> Dict[str, str]: """Get current status of all agents""" return self.agent_status.copy() # Convenience functions async def execute_red_team_scenario(scenario_name: str = "apt_simulation") -> Dict[str, Any]: """Execute a predefined red team scenario""" orchestrator = MultiAgentOrchestrator() template = orchestrator.get_scenario_template(scenario_name) if not template: raise ValueError(f"Unknown scenario template: {scenario_name}") return await orchestrator.execute_scenario(template) async def execute_vulnerability_assessment(targets: List[str]) -> Dict[str, Any]: """Execute vulnerability assessment scenario""" orchestrator = MultiAgentOrchestrator() template = orchestrator.get_scenario_template("vuln_assessment") if not template: raise ValueError("Vulnerability assessment template not found") # Customize template with specific targets context = {"custom_targets": targets} return await orchestrator.execute_scenario(template, context) # Example usage if __name__ == "__main__": async def main(): # Initialize orchestrator orchestrator = MultiAgentOrchestrator() # List available templates templates = orchestrator.list_scenario_templates() print(f"Available scenario templates: {templates}") # Execute APT simulation result = await execute_red_team_scenario("apt_simulation") print(f"Scenario completed: {result['evaluation']['overall_success']}") print(f"Success rate: {result['evaluation']['score']:.2%}") asyncio.run(main())