#!/usr/bin/env python3
"""
Exploitation Agent - Vulnerability exploitation and access gaining
"""

import json
import logging
from typing import Dict, List
from core.llm_manager import LLMManager
from tools.exploitation import (
    ExploitDatabase,
    MetasploitWrapper,
    WebExploiter,
    SQLInjector,
    RCEExploiter,
    BufferOverflowExploiter
)

logger = logging.getLogger(__name__)


class ExploitationAgent:
    """Agent responsible for vulnerability exploitation"""
    
    def __init__(self, config: Dict):
        """Initialize exploitation agent"""
        self.config = config
        self.llm = LLMManager(config)
        self.exploit_db = ExploitDatabase(config)
        self.metasploit = MetasploitWrapper(config)
        self.web_exploiter = WebExploiter(config)
        self.sql_injector = SQLInjector(config)
        self.rce_exploiter = RCEExploiter(config)
        self.bof_exploiter = BufferOverflowExploiter(config)
        
        logger.info("ExploitationAgent initialized")
    
    def execute(self, target: str, context: Dict) -> Dict:
        """Execute exploitation phase"""
        logger.info(f"Starting exploitation on {target}")
        
        results = {
            "target": target,
            "status": "running",
            "successful_exploits": [],
            "failed_attempts": [],
            "shells_obtained": [],
            "credentials_found": [],
            "ai_recommendations": {}
        }
        
        try:
            # Get reconnaissance data from context
            recon_data = context.get("phases", {}).get("recon", {})
            
            # Phase 1: Vulnerability Analysis
            logger.info("Phase 1: Analyzing vulnerabilities")
            vulnerabilities = self._identify_vulnerabilities(recon_data)
            
            # Phase 2: AI-powered Exploit Selection
            logger.info("Phase 2: AI exploit selection")
            exploit_plan = self._ai_exploit_planning(vulnerabilities, recon_data)
            results["ai_recommendations"] = exploit_plan
            
            # Phase 3: Execute Exploits
            logger.info("Phase 3: Executing exploits")
            for vuln in vulnerabilities[:5]:  # Limit to top 5 vulnerabilities
                exploit_result = self._attempt_exploitation(vuln, target)
                
                if exploit_result.get("success"):
                    results["successful_exploits"].append(exploit_result)
                    logger.info(f"Successful exploit: {vuln.get('type')}")
                    
                    # Check for shell access
                    if exploit_result.get("shell_access"):
                        results["shells_obtained"].append(exploit_result["shell_info"])
                else:
                    results["failed_attempts"].append(exploit_result)
            
            # Phase 4: Post-Exploitation Intelligence
            if results["successful_exploits"]:
                logger.info("Phase 4: Post-exploitation intelligence gathering")
                results["post_exploit_intel"] = self._gather_post_exploit_intel(
                    results["successful_exploits"]
                )
            
            results["status"] = "completed"
            logger.info("Exploitation phase completed")
            
        except Exception as e:
            logger.error(f"Error during exploitation: {e}")
            results["status"] = "error"
            results["error"] = str(e)
        
        return results
    
    def _identify_vulnerabilities(self, recon_data: Dict) -> List[Dict]:
        """Identify exploitable vulnerabilities from recon data"""
        vulnerabilities = []
        
        # Check network scan results
        network_scan = recon_data.get("network_scan", {})
        for host, data in network_scan.get("hosts", {}).items():
            for port in data.get("open_ports", []):
                vuln = {
                    "type": "network_service",
                    "host": host,
                    "port": port.get("port"),
                    "service": port.get("service"),
                    "version": port.get("version")
                }
                vulnerabilities.append(vuln)
        
        # Check web vulnerabilities
        web_analysis = recon_data.get("web_analysis", {})
        for vuln_type in ["sql_injection", "xss", "lfi", "rfi", "rce"]:
            if web_analysis.get(vuln_type):
                vulnerabilities.append({
                    "type": vuln_type,
                    "details": web_analysis[vuln_type]
                })
        
        return vulnerabilities
    
    def _ai_exploit_planning(self, vulnerabilities: List[Dict], recon_data: Dict) -> Dict:
        """Use AI to plan exploitation strategy"""
        prompt = self.llm.get_prompt(
            "exploitation",
            "ai_exploit_planning_user",
            default=f"""
Plan an exploitation strategy based on the following data:

Vulnerabilities Identified:
{json.dumps(vulnerabilities, indent=2)}

Reconnaissance Data:
{json.dumps(recon_data, indent=2)}

Provide:
1. Prioritized exploitation order
2. Recommended exploits for each vulnerability
3. Payload suggestions
4. Evasion techniques
5. Fallback strategies
6. Success probability estimates

Response in JSON format with detailed exploitation roadmap.
"""
        )
        
        system_prompt = self.llm.get_prompt(
            "exploitation",
            "ai_exploit_planning_system",
            default="""You are an expert exploit developer and penetration tester.
Create sophisticated exploitation plans considering detection, success rates, and impact.
Prioritize stealthy, reliable exploits over noisy attempts."""
        )
        
        try:
            formatted_prompt = prompt.format(
                vulnerabilities_json=json.dumps(vulnerabilities, indent=2),
                recon_data_json=json.dumps(recon_data, indent=2)
            )
            response = self.llm.generate(formatted_prompt, system_prompt)
            return json.loads(response)
        except Exception as e:
            logger.error(f"AI exploit planning error: {e}")
            return {"error": str(e)}
    
    def _attempt_exploitation(self, vulnerability: Dict, target: str) -> Dict:
        """Attempt to exploit a specific vulnerability"""
        vuln_type = vulnerability.get("type")
        
        result = {
            "vulnerability": vulnerability,
            "success": False,
            "method": None,
            "details": {}
        }
        
        try:
            if vuln_type == "sql_injection":
                result = self.sql_injector.exploit(target, vulnerability)
            elif vuln_type in ["xss", "csrf"]:
                result = self.web_exploiter.exploit(target, vulnerability)
            elif vuln_type in ["rce", "command_injection"]:
                result = self.rce_exploiter.exploit(target, vulnerability)
            elif vuln_type == "buffer_overflow":
                result = self.bof_exploiter.exploit(target, vulnerability)
            elif vuln_type == "network_service":
                result = self._exploit_network_service(target, vulnerability)
            else:
                # Use Metasploit for generic exploitation
                result = self.metasploit.exploit(target, vulnerability)
        
        except Exception as e:
            logger.error(f"Exploitation error for {vuln_type}: {e}")
            result["error"] = str(e)
        
        return result
    
    def _exploit_network_service(self, target: str, vulnerability: Dict) -> Dict:
        """Exploit network service vulnerabilities"""
        service = vulnerability.get("service", "").lower()
        
        # Check exploit database for known exploits
        exploits = self.exploit_db.search(service, vulnerability.get("version"))
        
        if exploits:
            logger.info(f"Found {len(exploits)} exploits for {service}")
            
            for exploit in exploits[:3]:  # Try top 3 exploits
                result = self.metasploit.run_exploit(
                    exploit["module"],
                    target,
                    vulnerability.get("port")
                )
                
                if result.get("success"):
                    return result
        
        return {"success": False, "message": "No suitable exploits found"}
    
    def _gather_post_exploit_intel(self, successful_exploits: List[Dict]) -> Dict:
        """Gather intelligence after successful exploitation"""
        intel = {
            "system_info": [],
            "user_accounts": [],
            "network_info": [],
            "installed_software": [],
            "credentials": []
        }
        
        for exploit in successful_exploits:
            if exploit.get("shell_access"):
                shell = exploit["shell_info"]
                
                # Gather system information
                # This would execute actual commands on compromised system
                # Placeholder for demonstration
                intel["system_info"].append({
                    "os": "detected_os",
                    "hostname": "detected_hostname",
                    "architecture": "x64"
                })
        
        return intel
    
    def generate_custom_exploit(self, vulnerability: Dict) -> str:
        """Generate custom exploit using AI"""
        target_info = {
            "vulnerability": vulnerability,
            "requirements": "Create working exploit code"
        }
        
        return self.llm.generate_payload(target_info, vulnerability.get("type"))
