"""
BREACH v3.0 - Phase 3: Exploitation
====================================

Parallel exploitation of vulnerability hypotheses.

NO EXPLOIT = NO REPORT

This phase:
1. Takes hypotheses from Phase 2
2. Attempts actual exploitation
3. Captures proof (screenshots, data, etc.)
4. Discards false positives
5. Returns only VALIDATED findings
"""

import asyncio
import re
import time
import hashlib
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional
from urllib.parse import quote

import aiohttp

from breach.ai import ClaudeAgent, HTTPTool, BrowserTool
from breach.ai.prompts import PromptManager
from .phase2_analysis import Hypothesis, AnalysisResult

try:
    from playwright.async_api import async_playwright
    PLAYWRIGHT_AVAILABLE = True
except ImportError:
    PLAYWRIGHT_AVAILABLE = False


@dataclass
class ValidatedFinding:
    """
    A validated, EXPLOITED vulnerability.

    This only exists if we successfully exploited the vulnerability.
    """
    id: str
    vuln_type: str
    severity: str  # CRITICAL, HIGH, MEDIUM, LOW
    endpoint: str
    parameter: str
    payload: str

    # PROOF - Required
    exploited: bool = True  # Always True for validated findings
    proof_type: str = ""  # sql_error, js_executed, data_extracted, etc.
    proof_data: Dict = field(default_factory=dict)
    evidence: str = ""

    # Screenshot (for browser-validated findings)
    screenshot: Optional[bytes] = None
    screenshot_hash: Optional[str] = None

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

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

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

    # Metadata
    discovered_at: datetime = field(default_factory=datetime.utcnow)
    exploitation_time_ms: float = 0


@dataclass
class ExploitResult:
    """Result of exploitation phase."""
    started_at: datetime = field(default_factory=datetime.utcnow)
    completed_at: Optional[datetime] = None
    duration_seconds: float = 0

    # Validated findings (EXPLOITED ONLY)
    findings: List[ValidatedFinding] = field(default_factory=list)

    # Stats
    hypotheses_tested: int = 0
    successful_exploits: int = 0
    false_positives: int = 0

    # By severity
    critical_count: int = 0
    high_count: int = 0
    medium_count: int = 0
    low_count: int = 0

    # Total business impact
    total_business_impact: int = 0

    def to_dict(self) -> Dict:
        return {
            "started_at": self.started_at.isoformat(),
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "duration_seconds": self.duration_seconds,
            "successful_exploits": self.successful_exploits,
            "false_positives": self.false_positives,
            "findings": [
                {
                    "id": f.id,
                    "vuln_type": f.vuln_type,
                    "severity": f.severity,
                    "endpoint": f.endpoint,
                    "parameter": f.parameter,
                    "proof_type": f.proof_type,
                    "curl_command": f.curl_command,
                }
                for f in self.findings
            ],
        }


class Exploiter:
    """Base exploiter for a vulnerability type."""

    VULN_TYPE: str = ""
    SEVERITY_MAP: Dict[str, str] = {}
    CWE_ID: str = ""
    REMEDIATION: str = ""

    # Business impact in USD
    IMPACT_HIGH: int = 50000
    IMPACT_MEDIUM: int = 10000
    IMPACT_LOW: int = 1000

    def __init__(
        self,
        session: aiohttp.ClientSession,
        browser_available: bool = False,
        audit_dir: Path = None,
    ):
        self.session = session
        self.browser_available = browser_available
        self.audit_dir = audit_dir
        self._browser = None
        self._page = None

    async def exploit(
        self,
        hypothesis: Hypothesis,
        cookies: Dict[str, str] = None,
    ) -> Optional[ValidatedFinding]:
        """Attempt to exploit the hypothesis. Returns finding if successful."""
        raise NotImplementedError

    async def _send_payload(
        self,
        url: str,
        param: str,
        payload: str,
        method: str = "GET",
        cookies: Dict = None,
    ) -> tuple:
        """Send payload and return (body, elapsed, status)."""
        try:
            start = time.time()
            if method.upper() == "GET":
                sep = "&" if "?" in url else "?"
                test_url = f"{url}{sep}{param}={quote(str(payload))}"
                async with self.session.get(test_url, cookies=cookies, ssl=False, timeout=15) as resp:
                    return await resp.text(), time.time() - start, resp.status
            else:
                async with self.session.post(url, data={param: payload}, cookies=cookies, ssl=False, timeout=15) as resp:
                    return await resp.text(), time.time() - start, resp.status
        except asyncio.TimeoutError:
            return "", 30, 0
        except:
            return "", 0, 0

    def _build_curl(self, url: str, param: str, payload: str, method: str = "GET") -> str:
        encoded = quote(str(payload))
        if method.upper() == "GET":
            sep = "&" if "?" in url else "?"
            return f"curl '{url}{sep}{param}={encoded}'"
        return f"curl -X POST '{url}' -d '{param}={encoded}'"


class SQLiExploiter(Exploiter):
    """SQL Injection exploiter."""

    VULN_TYPE = "sqli"
    CWE_ID = "CWE-89"
    REMEDIATION = "Use parameterized queries or prepared statements. Never concatenate user input into SQL."

    SQL_ERRORS = [
        (r"you have an error in your sql syntax", "MySQL"),
        (r"mysql.*error", "MySQL"),
        (r"postgres.*error", "PostgreSQL"),
        (r"pg_.*error", "PostgreSQL"),
        (r"sqlite.*error", "SQLite"),
        (r"ora-\d{5}", "Oracle"),
        (r"unclosed quotation mark", "MSSQL"),
    ]

    async def exploit(
        self,
        hypothesis: Hypothesis,
        cookies: Dict[str, str] = None,
    ) -> Optional[ValidatedFinding]:
        start_time = time.time()

        url = hypothesis.endpoint
        param = hypothesis.parameter
        payload = hypothesis.payload

        # Get baseline
        baseline, _, _ = await self._send_payload(url, param, "safe_test_123", "GET", cookies)

        # Send exploit payload
        body, elapsed, status = await self._send_payload(url, param, payload, "GET", cookies)

        # Check for SQL errors (not in baseline)
        for pattern, db_type in self.SQL_ERRORS:
            if re.search(pattern, body, re.IGNORECASE):
                if not re.search(pattern, baseline, re.IGNORECASE):
                    # EXPLOITED!
                    return ValidatedFinding(
                        id=f"sqli_{hashlib.sha256(f'{url}{param}{payload}'.encode()).hexdigest()[:8]}",
                        vuln_type=self.VULN_TYPE,
                        severity="CRITICAL",
                        endpoint=url,
                        parameter=param,
                        payload=payload,
                        proof_type="sql_error",
                        proof_data={"database": db_type, "error_pattern": pattern},
                        evidence=body[:500],
                        curl_command=self._build_curl(url, param, payload),
                        reproduction_steps=[
                            f"1. Send payload: {payload}",
                            f"2. Observe {db_type} error in response",
                        ],
                        business_impact=self.IMPACT_HIGH,
                        impact_explanation=f"SQL injection allows database access. Potential full database compromise.",
                        remediation=self.REMEDIATION,
                        cwe_id=self.CWE_ID,
                        exploitation_time_ms=(time.time() - start_time) * 1000,
                    )

        # Check for time-based
        if elapsed > 4.5 and any(t in payload.lower() for t in ["sleep", "waitfor"]):
            # Verify with baseline
            baseline_time = 0
            for _ in range(2):
                _, bt, _ = await self._send_payload(url, param, "safe", "GET", cookies)
                baseline_time += bt
            baseline_time /= 2

            if baseline_time < 2:
                return ValidatedFinding(
                    id=f"sqli_time_{hashlib.sha256(f'{url}{param}'.encode()).hexdigest()[:8]}",
                    vuln_type=self.VULN_TYPE,
                    severity="HIGH",
                    endpoint=url,
                    parameter=param,
                    payload=payload,
                    proof_type="time_delay",
                    proof_data={"delay_seconds": elapsed, "baseline_seconds": baseline_time},
                    evidence=f"Time delay: {elapsed:.2f}s vs baseline {baseline_time:.2f}s",
                    curl_command=self._build_curl(url, param, payload),
                    business_impact=self.IMPACT_HIGH,
                    remediation=self.REMEDIATION,
                    cwe_id=self.CWE_ID,
                    exploitation_time_ms=(time.time() - start_time) * 1000,
                )

        return None  # Not exploited


class XSSExploiter(Exploiter):
    """XSS exploiter - requires browser for proper validation."""

    VULN_TYPE = "xss"
    CWE_ID = "CWE-79"
    REMEDIATION = "Encode all user input before rendering in HTML. Use Content-Security-Policy header."

    async def exploit(
        self,
        hypothesis: Hypothesis,
        cookies: Dict[str, str] = None,
    ) -> Optional[ValidatedFinding]:
        start_time = time.time()

        if not self.browser_available or not PLAYWRIGHT_AVAILABLE:
            # XSS REQUIRES browser - cannot validate without it
            return None

        url = hypothesis.endpoint
        param = hypothesis.parameter
        payload = hypothesis.payload

        try:
            pw = await async_playwright().start()
            browser = await pw.chromium.launch(headless=True)
            page = await browser.new_page()

            # Track dialogs (alerts)
            dialog_triggered = []

            async def handle_dialog(dialog):
                dialog_triggered.append({"type": dialog.type, "message": dialog.message})
                await dialog.dismiss()

            page.on("dialog", handle_dialog)

            # Navigate with payload
            sep = "&" if "?" in url else "?"
            exploit_url = f"{url}{sep}{param}={quote(payload)}"

            await page.goto(exploit_url, wait_until="networkidle", timeout=15000)
            await page.wait_for_timeout(2000)

            # Check if XSS executed
            xss_confirmed = False
            proof_type = ""
            screenshot = None

            if dialog_triggered:
                xss_confirmed = True
                proof_type = "js_alert_executed"

            # Check DOM for injected script
            if not xss_confirmed:
                dom = await page.content()
                if payload in dom:
                    # Check if not escaped
                    escaped = ["&lt;script", "&lt;img", "\\u003c"]
                    if not any(e in dom for e in escaped):
                        xss_confirmed = True
                        proof_type = "dom_injection"

            # Capture screenshot as proof
            if xss_confirmed:
                screenshot = await page.screenshot(full_page=True)

            await browser.close()
            await pw.stop()

            if xss_confirmed:
                return ValidatedFinding(
                    id=f"xss_{hashlib.sha256(f'{url}{param}'.encode()).hexdigest()[:8]}",
                    vuln_type=self.VULN_TYPE,
                    severity="HIGH",
                    endpoint=url,
                    parameter=param,
                    payload=payload,
                    proof_type=proof_type,
                    proof_data={"dialogs": dialog_triggered},
                    evidence=f"JavaScript executed: {proof_type}",
                    screenshot=screenshot,
                    screenshot_hash=hashlib.sha256(screenshot).hexdigest() if screenshot else None,
                    curl_command=self._build_curl(url, param, payload),
                    reproduction_steps=[
                        f"1. Navigate to: {exploit_url}",
                        f"2. Observe JavaScript execution",
                    ],
                    business_impact=self.IMPACT_MEDIUM,
                    impact_explanation="XSS allows session hijacking and phishing attacks",
                    remediation=self.REMEDIATION,
                    cwe_id=self.CWE_ID,
                    exploitation_time_ms=(time.time() - start_time) * 1000,
                )

        except Exception as e:
            pass

        return None


class SSRFExploiter(Exploiter):
    """SSRF exploiter."""

    VULN_TYPE = "ssrf"
    CWE_ID = "CWE-918"
    REMEDIATION = "Validate and whitelist allowed URLs. Block requests to internal IPs and cloud metadata endpoints."

    METADATA_PATTERNS = [
        (r'ami-[a-z0-9]+', "AWS AMI ID"),
        (r'i-[a-z0-9]{8,17}', "AWS Instance ID"),
        (r'AKIA[A-Z0-9]{16}', "AWS Access Key"),
        (r'computeMetadata', "GCP Metadata"),
    ]

    async def exploit(
        self,
        hypothesis: Hypothesis,
        cookies: Dict[str, str] = None,
    ) -> Optional[ValidatedFinding]:
        start_time = time.time()

        url = hypothesis.endpoint
        param = hypothesis.parameter
        payload = hypothesis.payload

        # Get baseline
        baseline, _, _ = await self._send_payload(url, param, "https://example.com", "GET", cookies)

        # Send SSRF payload
        body, _, status = await self._send_payload(url, param, payload, "GET", cookies)

        # Check for cloud metadata
        for pattern, indicator in self.METADATA_PATTERNS:
            if re.search(pattern, body) and not re.search(pattern, baseline):
                return ValidatedFinding(
                    id=f"ssrf_{hashlib.sha256(f'{url}{param}'.encode()).hexdigest()[:8]}",
                    vuln_type=self.VULN_TYPE,
                    severity="CRITICAL",
                    endpoint=url,
                    parameter=param,
                    payload=payload,
                    proof_type="cloud_metadata_access",
                    proof_data={"indicator": indicator},
                    evidence=f"Cloud metadata accessed: {indicator}",
                    curl_command=self._build_curl(url, param, payload),
                    business_impact=self.IMPACT_HIGH,
                    impact_explanation="SSRF allows access to cloud credentials and internal services",
                    remediation=self.REMEDIATION,
                    cwe_id=self.CWE_ID,
                    exploitation_time_ms=(time.time() - start_time) * 1000,
                )

        return None


class CMDiExploiter(Exploiter):
    """Command injection exploiter."""

    VULN_TYPE = "cmdi"
    CWE_ID = "CWE-78"
    REMEDIATION = "Never pass user input to system commands. Use safe APIs instead of shell execution."

    CMD_PATTERNS = [
        (r'uid=\d+\([^)]+\)\s+gid=\d+', "id command"),
        (r'root:x?:0:0:', "/etc/passwd"),
        (r'Linux \w+ \d+\.\d+', "uname output"),
    ]

    async def exploit(
        self,
        hypothesis: Hypothesis,
        cookies: Dict[str, str] = None,
    ) -> Optional[ValidatedFinding]:
        start_time = time.time()

        url = hypothesis.endpoint
        param = hypothesis.parameter
        payload = hypothesis.payload

        baseline, _, _ = await self._send_payload(url, param, "safe", "GET", cookies)
        body, _, status = await self._send_payload(url, param, payload, "GET", cookies)

        for pattern, desc in self.CMD_PATTERNS:
            if re.search(pattern, body) and not re.search(pattern, baseline):
                return ValidatedFinding(
                    id=f"cmdi_{hashlib.sha256(f'{url}{param}'.encode()).hexdigest()[:8]}",
                    vuln_type=self.VULN_TYPE,
                    severity="CRITICAL",
                    endpoint=url,
                    parameter=param,
                    payload=payload,
                    proof_type="command_output",
                    proof_data={"output_type": desc},
                    evidence=f"Command executed: {desc}",
                    curl_command=self._build_curl(url, param, payload),
                    business_impact=self.IMPACT_HIGH,
                    impact_explanation="Command injection allows full system compromise",
                    remediation=self.REMEDIATION,
                    cwe_id=self.CWE_ID,
                    exploitation_time_ms=(time.time() - start_time) * 1000,
                )

        return None


class ExploitPhase:
    """
    Phase 3: Exploitation.

    Tests hypotheses through ACTUAL EXPLOITATION.
    Only returns findings that were successfully exploited.
    """

    EXPLOITERS = {
        "sqli": SQLiExploiter,
        "xss": XSSExploiter,
        "ssrf": SSRFExploiter,
        "cmdi": CMDiExploiter,
    }

    def __init__(
        self,
        use_browser: bool = True,
        max_concurrent: int = 5,
        audit_dir: Path = None,
        on_progress: Callable[[str], None] = None,
        on_finding: Callable[[ValidatedFinding], None] = None,
    ):
        self.use_browser = use_browser and PLAYWRIGHT_AVAILABLE
        self.max_concurrent = max_concurrent
        self.audit_dir = audit_dir
        self.on_progress = on_progress or (lambda x: None)
        self.on_finding = on_finding or (lambda x: None)

        self._session: Optional[aiohttp.ClientSession] = None

    async def __aenter__(self):
        self._session = aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=30),
            connector=aiohttp.TCPConnector(limit=self.max_concurrent, ssl=False),
        )
        return self

    async def __aexit__(self, *args):
        if self._session:
            await self._session.close()

    async def run(
        self,
        analysis_result: AnalysisResult,
        cookies: Dict[str, str] = None,
    ) -> ExploitResult:
        """
        Run exploitation on hypotheses.

        Args:
            analysis_result: Hypotheses from Phase 2
            cookies: Session cookies

        Returns:
            ExploitResult with ONLY validated (exploited) findings
        """
        start_time = time.time()
        result = ExploitResult()

        self.on_progress(f"Testing {len(analysis_result.hypotheses)} hypotheses...")

        if not self.use_browser:
            self.on_progress("[WARNING] Browser not available - XSS validation will be skipped")

        # Create exploiters
        exploiters = {
            vuln_type: ExploiterClass(
                session=self._session,
                browser_available=self.use_browser,
                audit_dir=self.audit_dir,
            )
            for vuln_type, ExploiterClass in self.EXPLOITERS.items()
        }

        # Create semaphore for concurrency control
        semaphore = asyncio.Semaphore(self.max_concurrent)

        async def try_exploit(hypothesis: Hypothesis) -> Optional[ValidatedFinding]:
            async with semaphore:
                exploiter = exploiters.get(hypothesis.vuln_type)
                if not exploiter:
                    return None

                try:
                    return await exploiter.exploit(hypothesis, cookies)
                except Exception as e:
                    return None

        # Run exploitations in parallel
        tasks = [try_exploit(h) for h in analysis_result.hypotheses]
        results = await asyncio.gather(*tasks)

        # Collect findings
        for hypothesis, finding in zip(analysis_result.hypotheses, results):
            result.hypotheses_tested += 1

            if finding:
                result.findings.append(finding)
                result.successful_exploits += 1
                self.on_finding(finding)
                self.on_progress(f"[EXPLOITED] {finding.vuln_type}: {finding.endpoint}")

                # Update severity counts
                if finding.severity == "CRITICAL":
                    result.critical_count += 1
                elif finding.severity == "HIGH":
                    result.high_count += 1
                elif finding.severity == "MEDIUM":
                    result.medium_count += 1
                else:
                    result.low_count += 1

                result.total_business_impact += finding.business_impact
            else:
                result.false_positives += 1

        # Finalize
        result.completed_at = datetime.utcnow()
        result.duration_seconds = time.time() - start_time

        self.on_progress(f"Exploitation complete: {result.successful_exploits} validated, {result.false_positives} false positives filtered")

        return result
