"""Attack Path Analyzer for MCP servers.

Analyzes MCP server configurations and source code to identify potential
attack paths — chains of entry points, intermediate nodes, and critical
assets that an attacker could exploit.

Builds an attack graph with known MCP attack patterns and generates
risk assessments with mitigation recommendations.

Zero external dependencies — pure Python stdlib.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


# ---------- Risk Levels ----------

class AttackRisk:
    """Attack path risk level constants."""
    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"


# ---------- Data Models ----------

@dataclass
class AttackNode:
    """A node in the attack graph."""

    node_id: str
    node_type: str  # entry_point, intermediate, critical_asset
    label: str
    description: str = ""
    file: str = ""
    line: int = 0

    def to_dict(self) -> dict[str, Any]:
        result: dict[str, Any] = {
            "node_id": self.node_id,
            "node_type": self.node_type,
            "label": self.label,
        }
        if self.description:
            result["description"] = self.description
        if self.file:
            result["file"] = self.file
        if self.line:
            result["line"] = self.line
        return result


@dataclass
class AttackEdge:
    """An edge in the attack graph connecting two nodes."""

    source: str  # node_id
    target: str  # node_id
    label: str = ""
    technique: str = ""

    def to_dict(self) -> dict[str, Any]:
        return {
            "source": self.source,
            "target": self.target,
            "label": self.label,
            "technique": self.technique,
        }


@dataclass
class AttackPath:
    """A single attack path from entry point to critical asset."""

    path_id: str
    source: str
    hops: list[str]
    target: str
    risk_level: str  # critical, high, medium, low
    description: str
    mitigation: str
    pattern: str = ""  # Named pattern (e.g., "user_input->tool_call->file_system")
    nodes: list[AttackNode] = field(default_factory=list)

    @property
    def length(self) -> int:
        return len(self.hops) + 2  # source + hops + target

    def to_dict(self) -> dict[str, Any]:
        return {
            "path_id": self.path_id,
            "source": self.source,
            "hops": self.hops,
            "target": self.target,
            "risk_level": self.risk_level,
            "description": self.description,
            "mitigation": self.mitigation,
            "pattern": self.pattern,
            "length": self.length,
            "nodes": [n.to_dict() for n in self.nodes],
        }


@dataclass
class AttackGraph:
    """Complete attack graph with nodes, edges, and paths."""

    nodes: list[AttackNode] = field(default_factory=list)
    edges: list[AttackEdge] = field(default_factory=list)
    paths: list[AttackPath] = field(default_factory=list)
    target_path: str = ""

    @property
    def total_paths(self) -> int:
        return len(self.paths)

    @property
    def critical_paths(self) -> int:
        return sum(1 for p in self.paths if p.risk_level == AttackRisk.CRITICAL)

    @property
    def high_paths(self) -> int:
        return sum(1 for p in self.paths if p.risk_level == AttackRisk.HIGH)

    @property
    def risk_summary(self) -> dict[str, int]:
        return {
            "critical": self.critical_paths,
            "high": self.high_paths,
            "medium": sum(1 for p in self.paths if p.risk_level == AttackRisk.MEDIUM),
            "low": sum(1 for p in self.paths if p.risk_level == AttackRisk.LOW),
        }

    @property
    def max_risk(self) -> str:
        if self.critical_paths > 0:
            return AttackRisk.CRITICAL
        if self.high_paths > 0:
            return AttackRisk.HIGH
        if any(p.risk_level == AttackRisk.MEDIUM for p in self.paths):
            return AttackRisk.MEDIUM
        if self.paths:
            return AttackRisk.LOW
        return "none"

    @property
    def exit_code(self) -> int:
        if self.critical_paths > 0:
            return 2
        if self.high_paths > 0:
            return 1
        return 0

    def to_dict(self) -> dict[str, Any]:
        return {
            "target_path": self.target_path,
            "total_paths": self.total_paths,
            "risk_summary": self.risk_summary,
            "max_risk": self.max_risk,
            "exit_code": self.exit_code,
            "nodes": [n.to_dict() for n in self.nodes],
            "edges": [e.to_dict() for e in self.edges],
            "paths": [p.to_dict() for p in self.paths],
        }


# ---------- Attack Patterns ----------

@dataclass
class AttackPattern:
    """A known attack pattern in MCP servers."""

    pattern_id: str
    name: str
    description: str
    chain: list[str]  # [entry_type, ..intermediate_types.., asset_type]
    risk_level: str
    mitigation: str


# Known MCP attack patterns
_PATTERNS: list[AttackPattern] = [
    AttackPattern(
        "ATCK-001",
        "user_input->tool_call->file_system",
        "User input flows through a tool call to file system operations, "
        "enabling path traversal or arbitrary file read/write",
        ["user_input", "tool_call", "file_system"],
        AttackRisk.CRITICAL,
        "Validate and sanitize all file paths. Use allowlists for permitted directories. "
        "Never pass raw user input to file operations.",
    ),
    AttackPattern(
        "ATCK-002",
        "user_input->prompt->llm->action",
        "User input is embedded in prompts sent to LLM, which then triggers "
        "actions — prompt injection leading to unauthorized actions",
        ["user_input", "prompt_construction", "llm_call", "action"],
        AttackRisk.CRITICAL,
        "Separate user input from system prompts. Use input/output guardrails. "
        "Validate LLM outputs before executing actions.",
    ),
    AttackPattern(
        "ATCK-003",
        "user_input->tool_call->command_execution",
        "User input reaches a tool that executes system commands, enabling "
        "remote code execution",
        ["user_input", "tool_call", "command_execution"],
        AttackRisk.CRITICAL,
        "Never pass user input to shell commands. Use parameterized APIs. "
        "Implement strict input validation and command allowlists.",
    ),
    AttackPattern(
        "ATCK-004",
        "user_input->tool_call->network_request",
        "User input controls URLs or parameters in network requests, "
        "enabling SSRF attacks against internal services",
        ["user_input", "tool_call", "network_request"],
        AttackRisk.HIGH,
        "Validate URLs against allowlists. Block internal/private IP ranges. "
        "Use a proxy with URL filtering.",
    ),
    AttackPattern(
        "ATCK-005",
        "user_input->tool_call->database_query",
        "User input flows into database queries, enabling SQL/NoSQL injection",
        ["user_input", "tool_call", "database_query"],
        AttackRisk.HIGH,
        "Use parameterized queries exclusively. Never concatenate user input "
        "into query strings. Use an ORM with proper escaping.",
    ),
    AttackPattern(
        "ATCK-006",
        "user_input->deserialization->code_execution",
        "User input is deserialized using unsafe methods, enabling "
        "arbitrary code execution",
        ["user_input", "deserialization", "code_execution"],
        AttackRisk.CRITICAL,
        "Never deserialize untrusted data. Use JSON instead of binary formats. "
        "If deserialization is needed, use safe loaders and validate schemas.",
    ),
    AttackPattern(
        "ATCK-007",
        "tool_call->env_var_access->secret_leak",
        "A tool accesses environment variables that contain secrets, "
        "which may be leaked to the user or logs",
        ["tool_call", "env_var_access", "secret_leak"],
        AttackRisk.HIGH,
        "Restrict environment variable access. Never expose secrets in tool responses. "
        "Use dedicated secret managers.",
    ),
    AttackPattern(
        "ATCK-008",
        "user_input->template->code_execution",
        "User input reaches a template engine, enabling Server-Side Template Injection",
        ["user_input", "template_engine", "code_execution"],
        AttackRisk.HIGH,
        "Never embed user input in templates. Use sandboxed template engines. "
        "Escape all dynamic content.",
    ),
    AttackPattern(
        "ATCK-009",
        "tool_escalation->permission_bypass",
        "One tool calls another with elevated permissions, bypassing access controls",
        ["tool_call", "tool_chain", "permission_bypass"],
        AttackRisk.HIGH,
        "Implement per-tool permission checks. Do not allow tools to call other tools "
        "without re-validating permissions.",
    ),
    AttackPattern(
        "ATCK-010",
        "user_input->resource_exhaustion",
        "User input triggers unbounded resource consumption — denial of service",
        ["user_input", "resource_consumption", "denial_of_service"],
        AttackRisk.MEDIUM,
        "Implement rate limiting, timeout, and resource quotas for all tools. "
        "Validate input size before processing.",
    ),
]


# ---------- Source Code Analysis ----------

# Regex patterns for detecting attack-relevant constructs in source code

_ENTRY_POINT_PATTERNS: list[tuple[str, str, re.Pattern[str]]] = [
    ("user_input", "Tool parameter (MCP tool input)",
     re.compile(r'(?:@tool|@server\.tool|def\s+\w+.*(?:request|input|query|params))', re.IGNORECASE)),
    ("user_input", "Request handler",
     re.compile(r'(?:@app\.route|@router|def\s+handle_|async\s+def\s+handle_)', re.IGNORECASE)),
    ("user_input", "Stdin/args input",
     re.compile(r'(?:sys\.argv|input\s*\(|stdin|argparse)', re.IGNORECASE)),
]

_INTERMEDIATE_PATTERNS: list[tuple[str, str, re.Pattern[str]]] = [
    ("tool_call", "Tool invocation",
     re.compile(r'(?:call_tool|invoke|execute_tool|run_tool)', re.IGNORECASE)),
    ("prompt_construction", "Prompt building",
     re.compile(r'(?:f["\'].*\{.*\}|\.format\(|prompt\s*[+=]|messages\.append)', re.IGNORECASE)),
    ("llm_call", "LLM API call",
     re.compile(r'(?:openai|anthropic|completion|chat\.create|generate\()', re.IGNORECASE)),
    ("template_engine", "Template rendering",
     re.compile(r'(?:render_template|Template\(|\.render\(|jinja)', re.IGNORECASE)),
    ("deserialization", "Unsafe deserialization",
     re.compile(r'(?:loads\(|yaml\.load|deserialize|fromstring)', re.IGNORECASE)),
    ("env_var_access", "Environment variable access",
     re.compile(r'(?:os\.environ|os\.getenv|environ\[|process\.env)', re.IGNORECASE)),
    ("tool_chain", "Tool-to-tool call",
     re.compile(r'(?:chain_tools|tool_pipeline|call_next|dispatch_tool)', re.IGNORECASE)),
    ("resource_consumption", "Resource-intensive operation",
     re.compile(r'(?:while\s+True|for.*range\(.*\)|sleep\(|open\()', re.IGNORECASE)),
]

_CRITICAL_ASSET_PATTERNS: list[tuple[str, str, re.Pattern[str]]] = [
    ("file_system", "File system access",
     re.compile(r'(?:open\(|\.read\(|\.write\(|Path\(|shutil\.|os\.remove|os\.unlink)', re.IGNORECASE)),
    ("command_execution", "Command execution",
     re.compile(r'(?:subprocess|os\.system|os\.popen|exec\(|eval\(|Popen)', re.IGNORECASE)),
    ("network_request", "Network request",
     re.compile(r'(?:requests\.|httpx\.|urllib|aiohttp|fetch\(|\.get\(|\.post\()', re.IGNORECASE)),
    ("database_query", "Database query",
     re.compile(r'(?:\.execute\(|\.query\(|cursor\.|\.find\(|\.aggregate\()', re.IGNORECASE)),
    ("code_execution", "Dynamic code execution",
     re.compile(r'(?:exec\(|eval\(|compile\(|__import__|importlib)', re.IGNORECASE)),
    ("secret_leak", "Secret/credential exposure",
     re.compile(r'(?:password|secret|token|api_key|credential|private_key)', re.IGNORECASE)),
    ("permission_bypass", "Permission check",
     re.compile(r'(?:check_permission|is_authorized|has_access|allowed)', re.IGNORECASE)),
    ("denial_of_service", "DoS target",
     re.compile(r'(?:timeout|rate_limit|max_retries|backoff)', re.IGNORECASE)),
]


def _scan_source_for_nodes(content: str, filepath: str) -> list[AttackNode]:
    """Scan source code and extract attack graph nodes."""
    nodes: list[AttackNode] = []
    lines = content.splitlines()
    node_counter = 0

    for line_num, line in enumerate(lines, start=1):
        for node_type, label, pattern in _ENTRY_POINT_PATTERNS:
            if pattern.search(line):
                node_counter += 1
                nodes.append(AttackNode(
                    node_id=f"N{node_counter:03d}",
                    node_type="entry_point",
                    label=label,
                    description=line.strip(),
                    file=filepath,
                    line=line_num,
                ))

        for node_type, label, pattern in _INTERMEDIATE_PATTERNS:
            if pattern.search(line):
                node_counter += 1
                nodes.append(AttackNode(
                    node_id=f"N{node_counter:03d}",
                    node_type="intermediate",
                    label=label,
                    description=line.strip(),
                    file=filepath,
                    line=line_num,
                ))

        for node_type, label, pattern in _CRITICAL_ASSET_PATTERNS:
            if pattern.search(line):
                node_counter += 1
                nodes.append(AttackNode(
                    node_id=f"N{node_counter:03d}",
                    node_type="critical_asset",
                    label=label,
                    description=line.strip(),
                    file=filepath,
                    line=line_num,
                ))

    return nodes


def _match_pattern_to_nodes(
    pattern: AttackPattern,
    nodes: list[AttackNode],
) -> list[AttackPath]:
    """Match an attack pattern against discovered nodes."""
    paths: list[AttackPath] = []
    chain = pattern.chain

    # Build type->nodes mapping
    type_map: dict[str, list[AttackNode]] = {}
    for node in nodes:
        label_lower = node.label.lower()
        for chain_type in chain:
            if chain_type in label_lower or _label_matches_type(label_lower, chain_type):
                type_map.setdefault(chain_type, []).append(node)

    # Check if we have nodes for each step in the chain
    if not all(chain_type in type_map for chain_type in chain):
        return paths

    # Generate paths: source -> hops -> target
    source_type = chain[0]
    target_type = chain[-1]
    hop_types = chain[1:-1]

    source_nodes = type_map.get(source_type, [])
    target_nodes = type_map.get(target_type, [])

    path_counter = 0
    for source_node in source_nodes:
        for target_node in target_nodes:
            if source_node.node_id == target_node.node_id:
                continue
            path_counter += 1
            hop_labels: list[str] = []
            path_nodes: list[AttackNode] = [source_node]
            for ht in hop_types:
                hop_nodes = type_map.get(ht, [])
                if hop_nodes:
                    hop_labels.append(hop_nodes[0].label)
                    path_nodes.append(hop_nodes[0])
            path_nodes.append(target_node)

            paths.append(AttackPath(
                path_id=f"{pattern.pattern_id}-P{path_counter:03d}",
                source=source_node.label,
                hops=hop_labels,
                target=target_node.label,
                risk_level=pattern.risk_level,
                description=pattern.description,
                mitigation=pattern.mitigation,
                pattern=pattern.name,
                nodes=path_nodes,
            ))
            if path_counter >= 10:
                return paths

    return paths


def _label_matches_type(label: str, chain_type: str) -> bool:
    """Check if a node label matches a chain type."""
    mappings: dict[str, list[str]] = {
        "user_input": ["tool parameter", "request handler", "stdin", "input"],
        "tool_call": ["tool invocation", "tool call", "tool parameter", "tool:"],
        "file_system": ["file system", "file access"],
        "command_execution": ["command execution", "dynamic code"],
        "network_request": ["network request", "http"],
        "database_query": ["database query", "database"],
        "code_execution": ["code execution", "dynamic code", "unsafe deseriali"],
        "prompt_construction": ["prompt build", "prompt"],
        "llm_call": ["llm api", "llm call"],
        "deserialization": ["deserializ", "unsafe deserializ"],
        "template_engine": ["template render", "template"],
        "env_var_access": ["environment variable"],
        "secret_leak": ["secret", "credential"],
        "tool_chain": ["tool-to-tool", "tool chain"],
        "permission_bypass": ["permission"],
        "resource_consumption": ["resource-intensive"],
        "denial_of_service": ["dos", "rate_limit", "timeout"],
        "action": ["action", "tool invocation", "command execution", "file system"],
    }
    match_list = mappings.get(chain_type, [])
    return any(m in label for m in match_list)


# ---------- Config-based Analysis ----------

def _analyze_config(config: dict[str, Any]) -> list[AttackNode]:
    """Extract attack nodes from MCP config JSON."""
    nodes: list[AttackNode] = []
    node_counter = 0

    for server_name, server_config in config.get("mcpServers", {}).items():
        if not isinstance(server_config, dict):
            continue

        cmd = server_config.get("command", "")
        if cmd:
            node_counter += 1
            nodes.append(AttackNode(
                node_id=f"C{node_counter:03d}",
                node_type="entry_point",
                label=f"Server: {server_name}",
                description=f"Command: {cmd}",
            ))

        tools = server_config.get("tools", [])
        if isinstance(tools, list):
            for tool in tools:
                name = tool if isinstance(tool, str) else tool.get("name", "unknown")
                node_counter += 1
                nodes.append(AttackNode(
                    node_id=f"C{node_counter:03d}",
                    node_type="intermediate",
                    label=f"Tool: {name}",
                    description=f"Tool in {server_name}",
                ))

        env = server_config.get("env", {})
        if isinstance(env, dict):
            for key in env:
                key_lower = key.lower()
                if any(s in key_lower for s in ("key", "secret", "token", "password", "auth")):
                    node_counter += 1
                    nodes.append(AttackNode(
                        node_id=f"C{node_counter:03d}",
                        node_type="critical_asset",
                        label=f"Secret: {key}",
                        description=f"Env var in {server_name}",
                    ))

        if "permissions" in server_config or "allowedDirectories" in server_config:
            node_counter += 1
            nodes.append(AttackNode(
                node_id=f"C{node_counter:03d}",
                node_type="critical_asset",
                label=f"Permissions: {server_name}",
                description="Access control configuration",
            ))

    return nodes


# ---------- Main Analyzer ----------

def analyze_attack_paths(path: Path) -> AttackGraph:
    """Analyze a file or directory for attack paths.

    Supports:
    - Python/JS/TS source files (source code analysis)
    - JSON config files (MCP config analysis)
    - Directories (recursive scan)

    Args:
        path: Path to file or directory.

    Returns:
        AttackGraph with all discovered paths.
    """
    graph = AttackGraph(target_path=str(path))

    if path.is_file():
        _analyze_single_file(path, graph)
    elif path.is_dir():
        _analyze_directory(path, graph)

    return graph


def _analyze_single_file(filepath: Path, graph: AttackGraph) -> None:
    """Analyze a single file and add results to graph."""
    try:
        content = filepath.read_text(encoding="utf-8", errors="replace")
    except (OSError, IOError):
        return

    suffix = filepath.suffix.lower()

    if suffix == ".json":
        _analyze_json_config(content, filepath, graph)
    elif suffix in (".py", ".js", ".ts", ".mjs", ".cjs"):
        _analyze_source_file(content, str(filepath), graph)


def _analyze_json_config(content: str, filepath: Path, graph: AttackGraph) -> None:
    """Analyze a JSON config file."""
    try:
        config = json.loads(content)
    except (json.JSONDecodeError, ValueError):
        return

    if not isinstance(config, dict):
        return

    config_nodes = _analyze_config(config)
    graph.nodes.extend(config_nodes)

    for ap_pattern in _PATTERNS:
        matched = _match_pattern_to_nodes(ap_pattern, config_nodes)
        graph.paths.extend(matched)

    _build_edges_from_paths(graph)


def _analyze_source_file(content: str, filepath: str, graph: AttackGraph) -> None:
    """Analyze a source code file."""
    nodes = _scan_source_for_nodes(content, filepath)
    graph.nodes.extend(nodes)

    for ap_pattern in _PATTERNS:
        matched = _match_pattern_to_nodes(ap_pattern, nodes)
        graph.paths.extend(matched)

    _build_edges_from_paths(graph)


def _analyze_directory(dirpath: Path, graph: AttackGraph) -> None:
    """Analyze all relevant files in a directory."""
    extensions = {".py", ".js", ".ts", ".mjs", ".cjs", ".json"}
    for filepath in sorted(dirpath.rglob("*")):
        if filepath.suffix.lower() not in extensions:
            continue
        parts = filepath.parts
        if "node_modules" in parts or ".git" in parts or "__pycache__" in parts:
            continue
        _analyze_single_file(filepath, graph)


def _build_edges_from_paths(graph: AttackGraph) -> None:
    """Build edges from discovered attack paths."""
    existing_edges: set[tuple[str, str]] = {(e.source, e.target) for e in graph.edges}
    for path in graph.paths:
        if not path.nodes:
            continue
        for i in range(len(path.nodes) - 1):
            src = path.nodes[i].node_id
            tgt = path.nodes[i + 1].node_id
            if (src, tgt) not in existing_edges:
                graph.edges.append(AttackEdge(
                    source=src,
                    target=tgt,
                    label=f"{path.nodes[i].label} -> {path.nodes[i + 1].label}",
                    technique=path.pattern,
                ))
                existing_edges.add((src, tgt))


# ---------- Formatters ----------

def format_attack_text(graph: AttackGraph) -> str:
    """Format an AttackGraph as human-readable text tree."""
    lines: list[str] = []
    lines.append(f"\n  Attack Path Analysis: {graph.target_path}")
    lines.append(f"  Total attack paths: {graph.total_paths}")
    lines.append(f"  Max risk: {graph.max_risk}")
    rs = graph.risk_summary
    lines.append(f"  Risk breakdown: critical={rs['critical']} "
                 f"high={rs['high']} medium={rs['medium']} low={rs['low']}")
    lines.append("")

    if not graph.paths:
        lines.append("  No attack paths detected.")
        return "\n".join(lines)

    risk_icons = {
        AttackRisk.CRITICAL: "[!!]",
        AttackRisk.HIGH: "[!]",
        AttackRisk.MEDIUM: "[-]",
        AttackRisk.LOW: "[.]",
    }

    for i, path in enumerate(graph.paths, 1):
        icon = risk_icons.get(path.risk_level, "[-]")
        lines.append(f"  {icon} Path {i}: {path.pattern}")
        lines.append(f"      Risk: {path.risk_level}")
        lines.append(f"      {path.source}")
        for hop in path.hops:
            lines.append(f"        -> {hop}")
        lines.append(f"        -> {path.target}")
        lines.append(f"      Mitigation: {path.mitigation[:80]}...")
        lines.append("")

    return "\n".join(lines)


def format_attack_json(graph: AttackGraph) -> str:
    """Format an AttackGraph as JSON."""
    return json.dumps(graph.to_dict(), indent=2, ensure_ascii=False)


def format_attack_dot(graph: AttackGraph) -> str:
    """Format an AttackGraph as DOT (Graphviz) format."""
    lines: list[str] = []
    lines.append("digraph attack_graph {")
    lines.append('  rankdir=LR;')
    lines.append('  node [shape=box, fontsize=10];')
    lines.append('')

    entry_nodes = [n for n in graph.nodes if n.node_type == "entry_point"]
    intermediate_nodes = [n for n in graph.nodes if n.node_type == "intermediate"]
    asset_nodes = [n for n in graph.nodes if n.node_type == "critical_asset"]

    if entry_nodes:
        lines.append('  // Entry points')
        for n in entry_nodes:
            lines.append(f'  "{n.node_id}" [label="{n.label}", style=filled, fillcolor="#90EE90"];')

    if intermediate_nodes:
        lines.append('  // Intermediate nodes')
        for n in intermediate_nodes:
            lines.append(f'  "{n.node_id}" [label="{n.label}", style=filled, fillcolor="#FFD700"];')

    if asset_nodes:
        lines.append('  // Critical assets')
        for n in asset_nodes:
            lines.append(f'  "{n.node_id}" [label="{n.label}", style=filled, fillcolor="#FF6347"];')

    lines.append('')
    lines.append('  // Edges')
    for edge in graph.edges:
        lines.append(f'  "{edge.source}" -> "{edge.target}" [label="{edge.label}"];')

    lines.append('}')
    return "\n".join(lines)


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

def get_attack_patterns() -> list[AttackPattern]:
    """Return all known attack patterns."""
    return list(_PATTERNS)
