"""
BREACH.AI - Evidence Collector
===============================

Comprehensive evidence collection for validated exploits.
Captures undeniable proof of successful exploitation.

Evidence Types:
- Screenshots (browser-based exploits)
- DOM snapshots (before/after exploitation)
- Network traces (requests/responses)
- Extracted data samples
- Console logs
- Video recordings (optional)
"""

import asyncio
import base64
import hashlib
import json
import re
import time
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple
from pathlib import Path

import aiohttp


@dataclass
class Evidence:
    """A piece of exploitation evidence."""
    evidence_type: str  # screenshot, dom, network, data, console, video
    description: str
    proves: str  # What does this prove?

    # Content
    content: Any  # Actual evidence content
    content_type: str  # image/png, text/html, application/json, etc.
    content_hash: str  # SHA256 for integrity

    # Metadata
    timestamp: datetime = field(default_factory=datetime.utcnow)
    file_path: Optional[str] = None

    # Severity
    severity: str = "INFO"  # CRITICAL, HIGH, MEDIUM, LOW, INFO

    # Redaction
    is_redacted: bool = False
    redaction_notes: Optional[str] = None


@dataclass
class EvidencePackage:
    """Complete evidence package for a finding."""
    finding_id: str
    vulnerability_type: str
    target: str

    # Evidence items
    evidence_items: List[Evidence] = field(default_factory=list)

    # Summary
    total_items: int = 0
    has_screenshot: bool = False
    has_data_sample: bool = False
    has_network_trace: bool = False

    # Timing
    collected_at: datetime = field(default_factory=datetime.utcnow)
    collection_time_ms: float = 0


class EvidenceCollector:
    """
    Collects and packages evidence for validated exploits.

    Captures comprehensive proof that can be used for:
    - Penetration test reports
    - Bug bounty submissions
    - Compliance audits
    - Incident response
    """

    def __init__(
        self,
        session: aiohttp.ClientSession,
        output_dir: str = "./evidence",
        capture_screenshots: bool = True,
        capture_network: bool = True,
        redact_sensitive: bool = True,
    ):
        self.session = session
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.capture_screenshots = capture_screenshots
        self.capture_network = capture_network
        self.redact_sensitive = redact_sensitive

    async def collect(
        self,
        finding: Any,
        validation_result: Any,
        exploit_result: Any = None,
    ) -> EvidencePackage:
        """
        Collect all evidence for a validated finding.

        Args:
            finding: The vulnerability finding
            validation_result: Result from exploitation validation
            exploit_result: Optional browser exploit result

        Returns:
            EvidencePackage with all collected evidence
        """
        start_time = time.time()

        package = EvidencePackage(
            finding_id=str(id(finding)),
            vulnerability_type=getattr(finding, 'category', 'unknown'),
            target=getattr(finding, 'endpoint', ''),
        )

        evidence_items = []

        # 1. Screenshot evidence
        if exploit_result and exploit_result.screenshot:
            evidence = self._create_screenshot_evidence(
                exploit_result.screenshot,
                finding,
            )
            evidence_items.append(evidence)
            package.has_screenshot = True

        # 2. DOM evidence
        if exploit_result and exploit_result.dom_after:
            evidence = self._create_dom_evidence(
                dom_before=exploit_result.dom_before,
                dom_after=exploit_result.dom_after,
                finding=finding,
            )
            evidence_items.append(evidence)

        # 3. Network trace evidence
        if exploit_result and exploit_result.network_requests:
            evidence = self._create_network_evidence(
                exploit_result.network_requests,
                finding,
            )
            evidence_items.append(evidence)
            package.has_network_trace = True

        # 4. Data extraction evidence
        if validation_result and validation_result.proof_data:
            evidence = self._create_data_evidence(
                validation_result.proof_data,
                finding,
            )
            evidence_items.append(evidence)
            package.has_data_sample = True

        # 5. Console log evidence
        if exploit_result and exploit_result.console_logs:
            evidence = self._create_console_evidence(
                exploit_result.console_logs,
                finding,
            )
            evidence_items.append(evidence)

        # 6. Request/Response evidence
        evidence = await self._create_request_response_evidence(
            finding,
            validation_result,
        )
        if evidence:
            evidence_items.append(evidence)

        # 7. Curl command evidence
        if validation_result and validation_result.curl_command:
            evidence = Evidence(
                evidence_type="reproduction",
                description="cURL command to reproduce the vulnerability",
                proves="Vulnerability can be reproduced with this command",
                content=validation_result.curl_command,
                content_type="text/plain",
                content_hash=hashlib.sha256(
                    validation_result.curl_command.encode()
                ).hexdigest(),
            )
            evidence_items.append(evidence)

        # Apply redaction if enabled
        if self.redact_sensitive:
            evidence_items = [self._redact_evidence(e) for e in evidence_items]

        package.evidence_items = evidence_items
        package.total_items = len(evidence_items)
        package.collection_time_ms = (time.time() - start_time) * 1000

        return package

    def _create_screenshot_evidence(
        self,
        screenshot: bytes,
        finding: Any,
    ) -> Evidence:
        """Create screenshot evidence."""
        # Save to file
        filename = f"screenshot_{int(time.time())}.png"
        filepath = self.output_dir / filename

        with open(filepath, "wb") as f:
            f.write(screenshot)

        return Evidence(
            evidence_type="screenshot",
            description=f"Screenshot of successful exploitation on {finding.endpoint}",
            proves="Visual proof of vulnerability exploitation",
            content=base64.b64encode(screenshot).decode()[:1000] + "...[truncated]",
            content_type="image/png",
            content_hash=hashlib.sha256(screenshot).hexdigest(),
            file_path=str(filepath),
            severity=getattr(finding, 'severity', 'HIGH'),
        )

    def _create_dom_evidence(
        self,
        dom_before: str,
        dom_after: str,
        finding: Any,
    ) -> Evidence:
        """Create DOM diff evidence."""
        # Create a simple diff
        diff = {
            "before_length": len(dom_before) if dom_before else 0,
            "after_length": len(dom_after) if dom_after else 0,
            "payload_present": finding.payload in (dom_after or ""),
            "dom_changed": dom_before != dom_after,
        }

        # Extract relevant sections
        if dom_after and finding.payload:
            # Find context around payload
            idx = dom_after.find(finding.payload)
            if idx != -1:
                start = max(0, idx - 200)
                end = min(len(dom_after), idx + len(finding.payload) + 200)
                diff["payload_context"] = dom_after[start:end]

        content = json.dumps(diff, indent=2)

        return Evidence(
            evidence_type="dom",
            description="DOM state comparison before and after exploitation",
            proves="Payload was injected into the DOM",
            content=content,
            content_type="application/json",
            content_hash=hashlib.sha256(content.encode()).hexdigest(),
        )

    def _create_network_evidence(
        self,
        requests: List[Dict],
        finding: Any,
    ) -> Evidence:
        """Create network trace evidence."""
        # Filter relevant requests
        relevant = []
        for req in requests[:20]:  # Limit to 20 requests
            relevant.append({
                "url": req.get("url", "")[:200],
                "method": req.get("method", ""),
                "status": req.get("status"),
            })

        content = json.dumps(relevant, indent=2)

        return Evidence(
            evidence_type="network",
            description=f"Network requests during exploitation ({len(relevant)} captured)",
            proves="Server communication during exploitation",
            content=content,
            content_type="application/json",
            content_hash=hashlib.sha256(content.encode()).hexdigest(),
        )

    def _create_data_evidence(
        self,
        proof_data: Dict,
        finding: Any,
    ) -> Evidence:
        """Create extracted data evidence."""
        # Redact sensitive values
        redacted = self._redact_dict(proof_data) if self.redact_sensitive else proof_data
        content = json.dumps(redacted, indent=2, default=str)

        return Evidence(
            evidence_type="data",
            description="Data extracted during exploitation",
            proves="Sensitive data was accessible/extractable",
            content=content,
            content_type="application/json",
            content_hash=hashlib.sha256(content.encode()).hexdigest(),
            severity="CRITICAL" if self._has_sensitive_data(proof_data) else "HIGH",
        )

    def _create_console_evidence(
        self,
        logs: List[str],
        finding: Any,
    ) -> Evidence:
        """Create console log evidence."""
        content = "\n".join(logs[:50])  # Limit to 50 logs

        return Evidence(
            evidence_type="console",
            description=f"Browser console output ({len(logs)} entries)",
            proves="JavaScript execution during exploitation",
            content=content,
            content_type="text/plain",
            content_hash=hashlib.sha256(content.encode()).hexdigest(),
        )

    async def _create_request_response_evidence(
        self,
        finding: Any,
        validation_result: Any,
    ) -> Optional[Evidence]:
        """Create HTTP request/response evidence."""
        if not validation_result:
            return None

        content = {
            "request": {
                "url": finding.endpoint,
                "method": getattr(finding, 'method', 'GET'),
                "parameter": finding.parameter,
                "payload": finding.payload,
            },
            "response": {
                "proof_type": validation_result.proof_type,
                "confidence": validation_result.confidence,
            },
        }

        content_str = json.dumps(content, indent=2)

        return Evidence(
            evidence_type="http",
            description="HTTP request and response details",
            proves="Vulnerability triggered by specific HTTP request",
            content=content_str,
            content_type="application/json",
            content_hash=hashlib.sha256(content_str.encode()).hexdigest(),
        )

    def _redact_evidence(self, evidence: Evidence) -> Evidence:
        """Redact sensitive information from evidence."""
        if evidence.content_type == "application/json":
            try:
                data = json.loads(evidence.content)
                redacted = self._redact_dict(data)
                evidence.content = json.dumps(redacted, indent=2)
                evidence.is_redacted = True
                evidence.redaction_notes = "Sensitive values partially redacted"
            except:
                pass
        elif evidence.content_type == "text/plain":
            evidence.content = self._redact_text(evidence.content)
            evidence.is_redacted = True

        return evidence

    def _redact_dict(self, data: Dict) -> Dict:
        """Redact sensitive values in dictionary."""
        sensitive_keys = [
            "password", "secret", "key", "token", "credential",
            "api_key", "access_key", "secret_key", "private",
            "ssn", "credit_card", "cvv",
        ]

        if isinstance(data, dict):
            redacted = {}
            for k, v in data.items():
                if any(s in k.lower() for s in sensitive_keys):
                    if isinstance(v, str) and len(v) > 8:
                        redacted[k] = v[:4] + "...[REDACTED]..." + v[-4:]
                    else:
                        redacted[k] = "[REDACTED]"
                elif isinstance(v, dict):
                    redacted[k] = self._redact_dict(v)
                elif isinstance(v, list):
                    redacted[k] = [self._redact_dict(i) if isinstance(i, dict) else i for i in v[:10]]
                else:
                    redacted[k] = v
            return redacted
        return data

    def _redact_text(self, text: str) -> str:
        """Redact sensitive patterns in text."""
        patterns = [
            # Email
            (r'([a-zA-Z0-9._%+-]+)@([a-zA-Z0-9.-]+\.[a-zA-Z]{2,})',
             r'\1[REDACTED]@\2'),
            # API keys
            (r'(AKIA|sk_live_|sk_test_)[A-Za-z0-9]{16,}',
             r'\1...[REDACTED]'),
            # Passwords in key=value
            (r'(password|secret|key|token)\s*[=:]\s*[^\s]+',
             r'\1=[REDACTED]'),
            # Credit card
            (r'\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b',
             r'****-****-****-****'),
        ]

        result = text
        for pattern, replacement in patterns:
            result = re.sub(pattern, replacement, result, flags=re.IGNORECASE)

        return result

    def _has_sensitive_data(self, data: Dict) -> bool:
        """Check if data contains sensitive information."""
        sensitive_indicators = [
            "password", "secret", "key", "token", "credential",
            "ssn", "credit_card", "email", "phone", "address",
        ]

        content = json.dumps(data).lower()
        return any(indicator in content for indicator in sensitive_indicators)

    def to_report(self, package: EvidencePackage) -> Dict:
        """Convert evidence package to report format."""
        return {
            "finding_id": package.finding_id,
            "vulnerability_type": package.vulnerability_type,
            "target": package.target,
            "evidence_summary": {
                "total_items": package.total_items,
                "has_screenshot": package.has_screenshot,
                "has_data_sample": package.has_data_sample,
                "has_network_trace": package.has_network_trace,
            },
            "evidence": [
                {
                    "type": e.evidence_type,
                    "description": e.description,
                    "proves": e.proves,
                    "severity": e.severity,
                    "hash": e.content_hash,
                    "content_preview": str(e.content)[:500] if e.content else None,
                    "file_path": e.file_path,
                    "is_redacted": e.is_redacted,
                }
                for e in package.evidence_items
            ],
            "collected_at": package.collected_at.isoformat(),
            "collection_time_ms": package.collection_time_ms,
        }
