"""
Proof-of-Exploit Engine — Feature 12 (v3.5.0).

For each HIGH finding in a supported vulnerability category, this engine:
  1. Selects safe PoC payloads for that category
  2. Spins up an ephemeral Docker container with the target code
  3. Sends payloads via HTTP and analyzes responses for exploitation signals
  4. Emits a three-tier verdict: verified-exploitable | verified-unexploitable | unverified

Safeguards:
  - Memory cap: --memory=128m
  - CPU cap: --cpus=0.5
  - Timeout: 30 seconds per finding
  - Port isolation: container exposed only on a random localhost port
  - Only operates on local target_dir code, never production URLs

This is DAST-augmented SAST: static analysis finds the candidate,
dynamic verification proves (or disproves) exploitability.
"""

from __future__ import annotations

import ast
import json
import logging
import os
import re
import socket
import subprocess
import time
import uuid
from dataclasses import asdict, dataclass, field
from pathlib import Path

import httpx

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

EXPLOIT_TIMEOUT = 30
CONTAINER_STARTUP_WAIT = 6  # seconds to wait for Flask to start
DEFAULT_APP_PORT = 5000

# Maps canonical rule IDs → exploit category
RULE_TO_CATEGORY: dict[str, str] = {
    # SQL injection
    "SECURITY-027": "sql-injection",
    "SECURITY-028": "sql-injection",
    # Command injection / shell execution
    "SECURITY-001": "command-injection",  # eval
    "SECURITY-021": "command-injection",  # shell=True
    "SECURITY-062": "command-injection",  # taint: req → exec
    # SSTI
    "SECURITY-032": "ssti",  # Mako templates
    # Path traversal
    "SECURITY-049": "path-traversal",
    "SECURITY-052": "path-traversal",
    # JS equivalents
    "SECURITY-053": "sql-injection",
    "SECURITY-054": "command-injection",
}

# Safe PoC payloads per category (cause observable behaviour without destructive side effects)
PAYLOADS: dict[str, list[str]] = {
    "sql-injection": [
        "1 OR 1=1--",
        "' OR '1'='1",
        "1 UNION SELECT sqlite_version()--",
        "'; SELECT sqlite_version(); --",
        "1 OR 1=1",
    ],
    "command-injection": [
        "localhost; echo EXPLOITED",
        "localhost && echo EXPLOITED",
        "localhost | echo EXPLOITED",
        "$(echo EXPLOITED)",
        "`echo EXPLOITED`",
    ],
    "ssti": [
        "{{7*7}}",
        "${7*7}",
        "#{7*7}",
        "{{config.items()|list}}",
    ],
    "path-traversal": [
        "../../../etc/hostname",
        "....//....//etc/hostname",
        "%2e%2e%2f%2e%2e%2f%2e%2e%2fetc%2fhostname",
        "../../../proc/version",
    ],
}

# Regex patterns that indicate successful exploitation
EXPLOITATION_SIGNALS: dict[str, list[str]] = {
    "sql-injection": [
        r"sqlite_version",
        r"near .+ syntax error",
        r"unrecognized token",
        r"1=1",
        r"\[\(1,\s*'admin'",  # row leak
        r"OperationalError",
        r"syntax error",
    ],
    "command-injection": [
        r"EXPLOITED",
        r"uid=\d+\(",
        r"root:x:0:0",
    ],
    "ssti": [
        r"\b49\b",  # {{7*7}} evaluated
        r"<class\s+",
        r"'DEBUG':",
        r"jinja2",
    ],
    "path-traversal": [
        r"root:x:0:0",
        r"nobody:x:",
        r"daemon:x:",
        r"Linux version",
        r"EXPLOITED",
    ],
}

# Common HTTP parameter names to try when we can't infer the param name
COMMON_PARAMS: dict[str, list[str]] = {
    "sql-injection": ["id", "user_id", "search", "q", "name", "filter"],
    "command-injection": ["host", "cmd", "command", "input", "arg", "exec"],
    "ssti": ["template", "name", "query", "msg", "text"],
    "path-traversal": ["path", "file", "filename", "page", "dir"],
}

# Default route patterns to try per category
DEFAULT_ROUTES: dict[str, list[str]] = {
    "sql-injection": ["/user", "/search", "/query", "/data", "/", "/api/user"],
    "command-injection": ["/ping", "/exec", "/run", "/cmd", "/"],
    "ssti": ["/render", "/template", "/page", "/"],
    "path-traversal": ["/file", "/read", "/download", "/static", "/"],
}


# ---------------------------------------------------------------------------
# Result dataclass
# ---------------------------------------------------------------------------


@dataclass
class ExploitResult:
    """Verdict from the exploit verifier for a single finding."""

    finding_id: str | int
    category: str
    verified: bool
    tier: str  # "verified-exploitable" | "verified-unexploitable" | "unverified"
    payload: str | None = None
    evidence: str | None = None
    container_id: str | None = None
    image_tag: str | None = None
    attempts: int = 0
    duration_seconds: float = 0.0
    error: str | None = None
    route: str | None = None
    param: str | None = None
    extra: dict = field(default_factory=dict)

    def to_dict(self) -> dict:
        return asdict(self)

    def to_proof_json(self) -> str:
        proof = {
            "payload": self.payload,
            "evidence": self.evidence,
            "container_id": self.container_id,
            "route": self.route,
            "param": self.param,
            "verified": self.verified,
            "tier": self.tier,
            "attempts": self.attempts,
            "duration_seconds": round(self.duration_seconds, 3),
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        }
        return json.dumps(proof)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _free_port() -> int:
    """Find a free localhost TCP port."""
    with socket.socket() as s:
        s.bind(("", 0))
        return s.getsockname()[1]


def _detect_exploitation(response: str, category: str) -> tuple[bool, str]:
    """Return (exploited, evidence_snippet) by matching category signals."""
    signals = EXPLOITATION_SIGNALS.get(category, [])
    for pattern in signals:
        m = re.search(pattern, response, re.IGNORECASE)
        if m:
            start = max(0, m.start() - 20)
            end = min(len(response), m.end() + 40)
            return True, response[start:end].strip()
    return False, ""


def _infer_route_and_param(file_path: str, line: int) -> tuple[str, str | None] | None:
    """
    Parse the file with AST to find the Flask route that contains `line`
    and infer the query-parameter name from `request.args.get(...)` calls.

    Returns (route_path, param_name) or None if not detectable.
    """
    try:
        source = Path(file_path).read_text()
        tree = ast.parse(source)
    except Exception:
        return None

    # Find the function that contains the target line
    target_func: ast.FunctionDef | ast.AsyncFunctionDef | None = None
    for node in ast.walk(tree):
        if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
            end = getattr(node, "end_lineno", None)
            if end is not None and node.lineno <= line <= end:
                if target_func is None:
                    target_func = node
                else:
                    prev_end = getattr(target_func, "end_lineno", end)
                    if (end - node.lineno) < (prev_end - target_func.lineno):
                        target_func = node

    if target_func is None:
        return None

    # Find the @app.route / @router.get decorator on that function
    route_path = None
    for decorator in target_func.decorator_list:
        route_path = _extract_route_from_decorator(decorator)
        if route_path:
            break

    # If no route on the function itself, look for the outermost function
    if route_path is None:
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
                for dec in node.decorator_list:
                    r = _extract_route_from_decorator(dec)
                    if r:
                        route_path = r
                        target_func = node
                        break
                if route_path:
                    break

    if route_path is None:
        return None

    # Find request.args.get('param') calls inside the function
    param_name: str | None = None
    for node in ast.walk(target_func):
        if isinstance(node, ast.Call):
            func = node.func
            if (
                isinstance(func, ast.Attribute)
                and func.attr == "get"
                and isinstance(func.value, ast.Attribute)
                and func.value.attr == "args"
            ):
                if node.args and isinstance(node.args[0], ast.Constant):
                    param_name = str(node.args[0].value)
                    break

    return (route_path, param_name)


def _extract_route_from_decorator(decorator: ast.expr) -> str | None:
    """Extract '/path' string from @app.route('/path') or @router.get('/path')."""
    if isinstance(decorator, ast.Call):
        args = decorator.args
        if args and isinstance(args[0], ast.Constant) and isinstance(args[0].value, str):
            val = args[0].value
            if val.startswith("/"):
                return val
    return None


def _write_default_dockerfile(target_dir: str) -> None:
    """Write a minimal Flask Dockerfile to target_dir if none exists."""
    df_path = Path(target_dir) / "Dockerfile"
    if df_path.exists():
        return
    df_path.write_text(
        "FROM python:3.11-slim\n"
        "WORKDIR /app\n"
        "COPY . .\n"
        "RUN pip install flask jinja2 --quiet --no-cache-dir\n"
        "EXPOSE 5000\n"
        'CMD ["python", "app.py"]\n'
    )


# ---------------------------------------------------------------------------
# Main engine
# ---------------------------------------------------------------------------


class ExploitVerifier:
    """
    Proves (or disproves) exploitability for HIGH findings in supported categories
    by running the target code inside an ephemeral Docker container.
    """

    def __init__(self, timeout: int = EXPLOIT_TIMEOUT, use_docker: bool = True):
        self.timeout = timeout
        self.use_docker = use_docker

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def is_docker_available(self) -> bool:
        """Return True if Docker daemon is reachable."""
        try:
            r = subprocess.run(
                ["docker", "info"],
                capture_output=True,
                timeout=5,
            )
            return r.returncode == 0
        except Exception:
            return False

    def can_verify(self, finding: dict) -> bool:
        """Return True if this finding is eligible for exploit verification."""
        severity = finding.get("severity", "").lower()
        if severity != "high":
            return False
        return self._infer_category(finding) is not None

    def verify_finding(self, finding: dict, target_dir: str) -> ExploitResult:
        """
        Attempt to prove or disprove exploitability of a single finding.

        Returns ExploitResult with tier:
          "verified-exploitable"   — payload succeeded, evidence captured
          "verified-unexploitable" — all payloads sent, none exploited
          "unverified"             — could not attempt (Docker unavailable, unsupported, etc.)
        """
        finding_id = finding.get("id", finding.get("canonical_rule_id", "unknown"))
        category = self._infer_category(finding)
        if category is None:
            return ExploitResult(
                finding_id=finding_id,
                category="unknown",
                verified=False,
                tier="unverified",
                error="Unsupported vulnerability category",
            )

        if not self.use_docker or not self.is_docker_available():
            return ExploitResult(
                finding_id=finding_id,
                category=category,
                verified=False,
                tier="unverified",
                error="Docker not available",
            )

        start = time.monotonic()
        image_tag = f"acrqa-exploit-{uuid.uuid4().hex[:8]}"
        container_id = None

        try:
            return self._run_verification(finding, target_dir, category, image_tag, start)
        except Exception as exc:
            logger.exception("ExploitVerifier.verify_finding failed")
            return ExploitResult(
                finding_id=finding_id,
                category=category,
                verified=False,
                tier="unverified",
                image_tag=image_tag,
                container_id=container_id,
                duration_seconds=time.monotonic() - start,
                error=str(exc),
            )
        finally:
            self._cleanup(container_id, image_tag)

    def enrich_findings(self, findings: list[dict], target_dir: str) -> list[dict]:
        """
        Attempt exploit verification for all eligible HIGH findings.

        Adds exploit_tier, exploit_proof, exploit_verified fields to each finding dict.
        Non-eligible findings are returned unchanged.
        Wrapped in try/except — never crashes the pipeline.
        """
        if not findings:
            return findings

        enriched = []
        for finding in findings:
            if not self.can_verify(finding):
                finding = dict(finding)
                finding.setdefault("exploit_tier", "unverified")
                enriched.append(finding)
                continue

            try:
                result = self.verify_finding(finding, target_dir)
                finding = dict(finding)
                finding["exploit_tier"] = result.tier
                finding["exploit_verified"] = result.verified
                finding["exploit_proof"] = result.to_proof_json() if result.verified else None
                if result.verified:
                    logger.info(
                        "ExploitVerifier: ✅ VERIFIED exploitable — %s (%s) payload=%r",
                        finding.get("canonical_rule_id", "?"),
                        result.category,
                        result.payload,
                    )
                else:
                    logger.debug(
                        "ExploitVerifier: finding %s tier=%s",
                        finding_id_str(finding),
                        result.tier,
                    )
            except Exception as _exc:
                logger.warning("ExploitVerifier: skipped finding %s: %s", finding_id_str(finding), _exc)
                finding = dict(finding)
                finding["exploit_tier"] = "unverified"
                finding["exploit_verified"] = False
                finding["exploit_proof"] = None

            enriched.append(finding)

        return enriched

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _infer_category(self, finding: dict) -> str | None:
        """Map finding to an exploit category, or None if unsupported."""
        rule_id = finding.get("canonical_rule_id") or finding.get("rule_id", "")
        # Direct rule → category mapping
        if rule_id in RULE_TO_CATEGORY:
            return RULE_TO_CATEGORY[rule_id]
        # Fall back to category field
        category = finding.get("category", "").lower()
        for cat in PAYLOADS:
            if cat in category:
                return cat
        # Check message for keywords
        message = finding.get("message", finding.get("description", "")).lower()
        for cat in PAYLOADS:
            keyword = cat.replace("-", " ")
            if keyword in message or cat in message:
                return cat
        return None

    def _run_verification(
        self,
        finding: dict,
        target_dir: str,
        category: str,
        image_tag: str,
        start: float,
    ) -> ExploitResult:
        finding_id = finding.get("id", finding.get("canonical_rule_id", "unknown"))

        # 1 — Detect route + param from AST
        file_path = finding.get("file", finding.get("file_path", ""))
        line = int(finding.get("line", finding.get("line_number", 0)) or 0)
        route_info = None
        if file_path and line:
            full_path = str(Path(target_dir) / file_path) if not os.path.isabs(file_path) else file_path
            if Path(full_path).exists():
                route_info = _infer_route_and_param(full_path, line)

        # 2 — Determine routes and params to try
        if route_info:
            routes = [route_info[0]]
            params = [route_info[1]] if route_info[1] else COMMON_PARAMS.get(category, ["q"])
        else:
            routes = DEFAULT_ROUTES.get(category, ["/"])
            params = COMMON_PARAMS.get(category, ["q"])

        payloads = PAYLOADS.get(category, [])

        # 3 — Build Docker image
        _write_default_dockerfile(target_dir)
        ok, build_err = self._build_image(target_dir, image_tag)
        if not ok:
            return ExploitResult(
                finding_id=finding_id,
                category=category,
                verified=False,
                tier="unverified",
                image_tag=image_tag,
                duration_seconds=time.monotonic() - start,
                error=f"Docker build failed: {build_err[:200]}",
            )

        # 4 — Start container
        port = _free_port()
        container_id, run_err = self._start_container(image_tag, port)
        if container_id is None:
            return ExploitResult(
                finding_id=finding_id,
                category=category,
                verified=False,
                tier="unverified",
                image_tag=image_tag,
                duration_seconds=time.monotonic() - start,
                error=f"Container start failed: {run_err[:200]}",
            )

        # 5 — Wait for app startup
        self._wait_for_http(port, timeout=self.timeout)

        # 6 — Send payloads
        attempts = 0
        for route in routes:
            for param in params:
                for payload in payloads:
                    if time.monotonic() - start > self.timeout:
                        break
                    attempts += 1
                    try:
                        response = self._send_payload(port, route, param, payload)
                        exploited, evidence = _detect_exploitation(response, category)
                        if exploited:
                            duration = time.monotonic() - start
                            self._cleanup(container_id, image_tag)
                            return ExploitResult(
                                finding_id=finding_id,
                                category=category,
                                verified=True,
                                tier="verified-exploitable",
                                payload=payload,
                                evidence=evidence,
                                container_id=container_id,
                                image_tag=image_tag,
                                attempts=attempts,
                                duration_seconds=duration,
                                route=route,
                                param=param,
                            )
                    except Exception as _req_err:
                        logger.debug("Payload attempt failed: %s", _req_err)

        duration = time.monotonic() - start
        self._cleanup(container_id, image_tag)
        return ExploitResult(
            finding_id=finding_id,
            category=category,
            verified=False,
            tier="verified-unexploitable",
            container_id=container_id,
            image_tag=image_tag,
            attempts=attempts,
            duration_seconds=duration,
        )

    def _build_image(self, target_dir: str, image_tag: str) -> tuple[bool, str]:
        """Build Docker image from target_dir. Returns (success, error_output)."""
        try:
            r = subprocess.run(
                ["docker", "build", "-t", image_tag, "."],
                cwd=target_dir,
                capture_output=True,
                text=True,
                timeout=self.timeout,
            )
            return r.returncode == 0, r.stderr
        except subprocess.TimeoutExpired:
            return False, "Docker build timed out"
        except Exception as exc:
            return False, str(exc)

    def _start_container(self, image_tag: str, port: int) -> tuple[str | None, str]:
        """Start container with resource caps. Returns (container_id, error)."""
        try:
            r = subprocess.run(
                [
                    "docker",
                    "run",
                    "-d",
                    "--rm",
                    "-p",
                    f"{port}:{DEFAULT_APP_PORT}",
                    "--memory=128m",
                    "--cpus=0.5",
                    image_tag,
                ],
                capture_output=True,
                text=True,
                timeout=15,
            )
            if r.returncode == 0:
                return r.stdout.strip(), ""
            return None, r.stderr
        except subprocess.TimeoutExpired:
            return None, "Container start timed out"
        except Exception as exc:
            return None, str(exc)

    def _wait_for_http(self, port: int, timeout: int = 10) -> bool:
        """Poll until HTTP server is ready or timeout."""
        deadline = time.time() + timeout
        while time.time() < deadline:
            try:
                with httpx.Client(timeout=1) as client:
                    client.get(f"http://localhost:{port}/health")
                    return True
            except Exception:
                time.sleep(0.5)
        return False

    def _send_payload(self, port: int, path: str, param: str, payload: str) -> str:
        """Send GET request with payload as query param, return response text."""
        url = f"http://localhost:{port}{path}"
        with httpx.Client(timeout=5) as client:
            r = client.get(url, params={param: payload})
            return r.text

    def _cleanup(self, container_id: str | None, image_tag: str) -> None:
        """Stop container and remove image (best-effort)."""
        if container_id:
            try:
                subprocess.run(
                    ["docker", "stop", container_id],
                    capture_output=True,
                    timeout=10,
                )
            except Exception:
                pass
        try:
            subprocess.run(
                ["docker", "rmi", "-f", image_tag],
                capture_output=True,
                timeout=10,
            )
        except Exception:
            pass


# ---------------------------------------------------------------------------
# Utility
# ---------------------------------------------------------------------------


def finding_id_str(finding: dict) -> str:
    return str(finding.get("id") or finding.get("canonical_rule_id") or "?")
