"""MCP Server Fingerprinting -- Detect framework and apply specific rules.

Analyzes source code to determine which MCP framework a server uses
(FastMCP, mcp-python-sdk, @modelcontextprotocol/sdk, custom) and
provides framework-specific security recommendations and rules.

Zero external dependencies. All detection is regex-based.
"""

from __future__ import annotations

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


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

@dataclass
class KnownVulnerability:
    """A known vulnerability for a specific MCP framework.

    Attributes:
        id: Vulnerability identifier.
        title: Short title.
        description: Detailed description.
        severity: Severity string (CRITICAL, HIGH, MEDIUM, LOW).
        cwe: CWE reference.
        affected_versions: Version range affected.
        fix: Remediation guidance.
    """

    id: str
    title: str
    description: str
    severity: str = "MEDIUM"
    cwe: str = ""
    affected_versions: str = ""
    fix: str = ""

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "title": self.title,
            "description": self.description,
            "severity": self.severity,
            "cwe": self.cwe,
            "affected_versions": self.affected_versions,
            "fix": self.fix,
        }


@dataclass
class FrameworkInfo:
    """Information about a detected MCP framework.

    Attributes:
        name: Framework name (e.g. 'FastMCP', 'mcp-python-sdk').
        version_hint: Detected or inferred version hint.
        known_vulnerabilities: Known security issues for this framework.
        recommended_rules: Rule IDs recommended for this framework.
        confidence: Detection confidence (HIGH, MEDIUM, LOW).
        details: Additional detection details.
    """

    name: str
    version_hint: str = ""
    known_vulnerabilities: list[KnownVulnerability] = field(default_factory=list)
    recommended_rules: list[str] = field(default_factory=list)
    confidence: str = "HIGH"
    details: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "version_hint": self.version_hint,
            "known_vulnerabilities": [v.to_dict() for v in self.known_vulnerabilities],
            "recommended_rules": self.recommended_rules,
            "confidence": self.confidence,
            "details": self.details,
        }


# ---------------------------------------------------------------------------
# Framework Definitions
# ---------------------------------------------------------------------------

# Known vulnerabilities per framework
_FASTMCP_VULNS: list[KnownVulnerability] = [
    KnownVulnerability(
        id="FASTMCP-001",
        title="Auto-exposure of all decorated functions",
        description=(
            "FastMCP auto-exposes all @tool decorated functions without "
            "explicit access control. This can expose internal functions "
            "that were not intended to be public."
        ),
        severity="HIGH",
        cwe="CWE-284",
        fix="Explicitly list tools to expose. Use tool access control decorators.",
    ),
    KnownVulnerability(
        id="FASTMCP-002",
        title="Default permissive CORS in development mode",
        description=(
            "FastMCP development mode enables permissive CORS by default, "
            "which can be accidentally deployed to production."
        ),
        severity="MEDIUM",
        cwe="CWE-942",
        fix="Set explicit CORS origins. Never use development mode in production.",
    ),
    KnownVulnerability(
        id="FASTMCP-003",
        title="Missing rate limiting by default",
        description=(
            "FastMCP does not include rate limiting by default. "
            "Tools are exposed without any request throttling."
        ),
        severity="MEDIUM",
        cwe="CWE-770",
        fix="Add rate limiting middleware or decorators to all exposed tools.",
    ),
]

_MCP_PYTHON_SDK_VULNS: list[KnownVulnerability] = [
    KnownVulnerability(
        id="MCP-SDK-001",
        title="stdio transport without authentication",
        description=(
            "mcp-python-sdk stdio transport does not include authentication. "
            "Any process that can connect via stdio gets full access."
        ),
        severity="MEDIUM",
        cwe="CWE-306",
        fix="Implement authentication layer on top of stdio transport.",
    ),
    KnownVulnerability(
        id="MCP-SDK-002",
        title="Tool descriptions used as prompt context",
        description=(
            "Tool descriptions in mcp-python-sdk are passed directly to the LLM. "
            "Malicious descriptions can influence model behavior (prompt injection)."
        ),
        severity="HIGH",
        cwe="CWE-74",
        fix="Sanitize tool descriptions. Avoid dynamic content in descriptions.",
    ),
]

_TS_SDK_VULNS: list[KnownVulnerability] = [
    KnownVulnerability(
        id="TS-SDK-001",
        title="No input validation in handler types",
        description=(
            "@modelcontextprotocol/sdk TypeScript handlers accept 'any' types "
            "by default. Without explicit Zod/io-ts validation, type safety "
            "is not enforced at runtime."
        ),
        severity="HIGH",
        cwe="CWE-20",
        fix="Use Zod schemas for all tool handler inputs.",
    ),
    KnownVulnerability(
        id="TS-SDK-002",
        title="Error messages may leak internal details",
        description=(
            "Default error handling in @modelcontextprotocol/sdk passes "
            "error messages directly to clients, potentially leaking "
            "internal paths and stack traces."
        ),
        severity="MEDIUM",
        cwe="CWE-209",
        fix="Implement custom error handler that sanitizes error messages.",
    ),
]

_CUSTOM_VULNS: list[KnownVulnerability] = [
    KnownVulnerability(
        id="CUSTOM-001",
        title="Custom MCP implementation without standard security controls",
        description=(
            "Custom MCP implementations lack the security features of "
            "established SDKs (input validation, error handling, transport "
            "security)."
        ),
        severity="HIGH",
        cwe="CWE-693",
        fix="Consider using an established MCP SDK. Implement all OWASP Agentic controls.",
    ),
]

# Recommended rules per framework
_FRAMEWORK_RULES: dict[str, list[str]] = {
    "FastMCP": [
        "PERM-001", "RATE-001", "MCP-001", "MCP-002", "AUTH-001",
        "GEN-MCP-001", "GEN-MCP-003", "GEN-MCP-005",
    ],
    "mcp-python-sdk": [
        "PATH-001", "CMD-001", "CMD-002", "MCP-001", "MCP-002",
        "MCP-003", "MCP-004", "GEN-MCP-001",
    ],
    "@modelcontextprotocol/sdk": [
        "TS-CMD-001", "TS-SQL-001", "TS-PATH-001", "TS-INPUT-001",
        "TS-AUTH-001", "TS-PROMPT-001",
    ],
    "custom": [
        "PATH-001", "CMD-001", "CMD-002", "AUTH-001", "INPUT-001",
        "RATE-001", "SSRF-001", "DESER-001",
    ],
}


# ---------------------------------------------------------------------------
# Detection Patterns
# ---------------------------------------------------------------------------

_FASTMCP_PATTERNS: list[re.Pattern[str]] = [
    re.compile(r"from\s+fastmcp\s+import", re.IGNORECASE),
    re.compile(r"import\s+fastmcp", re.IGNORECASE),
    re.compile(r"FastMCP\s*\("),
    re.compile(r"@\w+\.tool\b"),
]

_MCP_PYTHON_PATTERNS: list[re.Pattern[str]] = [
    re.compile(r"from\s+mcp\.server\s+import"),
    re.compile(r"from\s+mcp\s+import\s+.*Server"),
    re.compile(r"stdio_server\s*\("),
    re.compile(r"mcp\.server\.Server\s*\("),
]

_TS_SDK_PATTERNS: list[re.Pattern[str]] = [
    re.compile(r"@modelcontextprotocol/sdk"),
    re.compile(r"from\s+['\"]@modelcontextprotocol"),
    re.compile(r"import\s+.*from\s+['\"]@modelcontextprotocol"),
    re.compile(r"McpServer\s*\("),
    re.compile(r"createServer\s*\("),
]

_VERSION_PATTERNS: dict[str, re.Pattern[str]] = {
    "fastmcp": re.compile(r"fastmcp[=<>~!]+([0-9]+\.[0-9]+(?:\.[0-9]+)?)"),
    "mcp": re.compile(r"mcp[=<>~!]+([0-9]+\.[0-9]+(?:\.[0-9]+)?)"),
    "mcp_sdk": re.compile(
        r"@modelcontextprotocol/sdk[\"']?\s*:\s*[\"']?\^?~?([0-9]+\.[0-9]+(?:\.[0-9]+)?)"
    ),
}


# ---------------------------------------------------------------------------
# Fingerprinting Functions
# ---------------------------------------------------------------------------

def fingerprint_server(source_code: str) -> FrameworkInfo:
    """Detect which MCP framework a server uses from its source code.

    Analyzes source code for import patterns, SDK usage, and framework
    markers to determine the MCP framework in use.

    Args:
        source_code: Source code string to analyze.

    Returns:
        FrameworkInfo with detected framework, vulnerabilities, and rules.
    """
    # Score each framework
    scores: dict[str, int] = {
        "FastMCP": 0,
        "mcp-python-sdk": 0,
        "@modelcontextprotocol/sdk": 0,
    }
    matched_patterns: dict[str, list[str]] = {k: [] for k in scores}

    for pat in _FASTMCP_PATTERNS:
        match = pat.search(source_code)
        if match:
            scores["FastMCP"] += 1
            matched_patterns["FastMCP"].append(match.group())

    for pat in _MCP_PYTHON_PATTERNS:
        match = pat.search(source_code)
        if match:
            scores["mcp-python-sdk"] += 1
            matched_patterns["mcp-python-sdk"].append(match.group())

    for pat in _TS_SDK_PATTERNS:
        match = pat.search(source_code)
        if match:
            scores["@modelcontextprotocol/sdk"] += 1
            matched_patterns["@modelcontextprotocol/sdk"].append(match.group())

    # Find winner
    max_score = max(scores.values())
    if max_score == 0:
        # No known framework detected
        return FrameworkInfo(
            name="custom",
            version_hint="",
            known_vulnerabilities=list(_CUSTOM_VULNS),
            recommended_rules=_FRAMEWORK_RULES["custom"],
            confidence="LOW",
            details={"reason": "No known MCP framework patterns detected"},
        )

    winner = max(scores, key=lambda k: scores[k])

    # Determine confidence
    if max_score >= 3:
        confidence = "HIGH"
    elif max_score >= 2:
        confidence = "MEDIUM"
    else:
        confidence = "LOW"

    # Detect version hint
    version_hint = _detect_version(source_code, winner)

    # Get known vulns and rules
    vulns_map = {
        "FastMCP": _FASTMCP_VULNS,
        "mcp-python-sdk": _MCP_PYTHON_SDK_VULNS,
        "@modelcontextprotocol/sdk": _TS_SDK_VULNS,
    }

    return FrameworkInfo(
        name=winner,
        version_hint=version_hint,
        known_vulnerabilities=list(vulns_map.get(winner, [])),
        recommended_rules=_FRAMEWORK_RULES.get(winner, []),
        confidence=confidence,
        details={
            "score": max_score,
            "matched_patterns": matched_patterns[winner],
        },
    )


def fingerprint_file(filepath: Path) -> FrameworkInfo:
    """Fingerprint an MCP server from a file path.

    Args:
        filepath: Path to the source file.

    Returns:
        FrameworkInfo for the detected framework.
    """
    try:
        source = filepath.read_text(encoding="utf-8", errors="replace")
    except (OSError, PermissionError):
        return FrameworkInfo(
            name="unknown",
            confidence="LOW",
            details={"error": f"Could not read file: {filepath}"},
        )
    return fingerprint_server(source)


def fingerprint_directory(dirpath: Path) -> list[tuple[Path, FrameworkInfo]]:
    """Fingerprint all source files in a directory.

    Args:
        dirpath: Directory to scan.

    Returns:
        List of (filepath, FrameworkInfo) tuples for detected servers.
    """
    results: list[tuple[Path, FrameworkInfo]] = []
    extensions = {".py", ".ts", ".js", ".tsx", ".jsx", ".mjs"}
    skip_dirs = {"node_modules", ".git", "__pycache__", ".venv", "venv"}

    if not dirpath.is_dir():
        return results

    for item in sorted(dirpath.rglob("*")):
        if item.is_file() and item.suffix in extensions:
            # Skip excluded dirs
            if any(part in skip_dirs for part in item.parts):
                continue
            info = fingerprint_file(item)
            if info.name != "unknown":
                results.append((item, info))

    return results


def _detect_version(source: str, framework: str) -> str:
    """Try to detect framework version from source or config patterns.

    Args:
        source: Source code to search.
        framework: Framework name to search version for.

    Returns:
        Version string if detected, empty string otherwise.
    """
    key_map = {
        "FastMCP": "fastmcp",
        "mcp-python-sdk": "mcp",
        "@modelcontextprotocol/sdk": "mcp_sdk",
    }
    key = key_map.get(framework)
    if key and key in _VERSION_PATTERNS:
        match = _VERSION_PATTERNS[key].search(source)
        if match:
            return match.group(1)
    return ""


# ---------------------------------------------------------------------------
# CLI Formatting
# ---------------------------------------------------------------------------

def format_fingerprint_text(info: FrameworkInfo, filepath: str = "") -> str:
    """Format fingerprint result as human-readable text.

    Args:
        info: FrameworkInfo to format.
        filepath: Optional filepath for context.

    Returns:
        Formatted text string.
    """
    lines: list[str] = []
    lines.append(f"\n{'=' * 60}")
    lines.append("  mcpaudit -- MCP Server Fingerprint")
    lines.append(f"{'=' * 60}")

    if filepath:
        lines.append(f"  File:       {filepath}")
    lines.append(f"  Framework:  {info.name}")
    if info.version_hint:
        lines.append(f"  Version:    {info.version_hint}")
    lines.append(f"  Confidence: {info.confidence}")

    if info.known_vulnerabilities:
        lines.append(f"\n  Known Vulnerabilities ({len(info.known_vulnerabilities)}):")
        for vuln in info.known_vulnerabilities:
            lines.append(f"    [{vuln.severity}] {vuln.id}: {vuln.title}")
            if vuln.fix:
                lines.append(f"           Fix: {vuln.fix}")

    if info.recommended_rules:
        lines.append(f"\n  Recommended Rules: {', '.join(info.recommended_rules)}")

    lines.append(f"\n{'=' * 60}\n")
    return "\n".join(lines)


def format_fingerprint_json(info: FrameworkInfo, filepath: str = "") -> str:
    """Format fingerprint result as JSON.

    Args:
        info: FrameworkInfo to format.
        filepath: Optional filepath for context.

    Returns:
        JSON string.
    """
    import json
    data = info.to_dict()
    if filepath:
        data["file"] = filepath
    return json.dumps(data, indent=2, ensure_ascii=False)


# ---------------------------------------------------------------------------
# MCPFingerprint — Extended fingerprint with transport, tools, resources
# ---------------------------------------------------------------------------

@dataclass
class MCPFingerprint:
    """Extended fingerprint for an MCP server file or directory.

    Combines framework detection with transport, tool/resource enumeration,
    and version hints for a comprehensive server profile.
    """

    framework: str = "unknown"
    language: str = "unknown"
    version_hints: list[str] = field(default_factory=list)
    transport: str = "unknown"
    capabilities: list[str] = field(default_factory=list)
    tool_count: int = 0
    resource_count: int = 0
    tools: list[str] = field(default_factory=list)
    resources: list[str] = field(default_factory=list)
    file_path: str = ""
    errors: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "framework": self.framework,
            "language": self.language,
            "version_hints": self.version_hints,
            "transport": self.transport,
            "capabilities": self.capabilities,
            "tool_count": self.tool_count,
            "resource_count": self.resource_count,
            "tools": self.tools,
            "resources": self.resources,
            "file_path": self.file_path,
        }


# Transport detection patterns
_STDIO_PATTERN = re.compile(
    r"stdio_server|StdioServerTransport|"
    r"transport.*stdio|type.*stdio",
    re.IGNORECASE,
)
_SSE_PATTERN = re.compile(
    r"sse_server|SSEServerTransport|SseServerTransport|"
    r"transport.*sse|type.*sse",
    re.IGNORECASE,
)
_HTTP_PATTERN = re.compile(
    r"streamable.http|StreamableHTTP|"
    r"transport.*http|type.*streamable",
    re.IGNORECASE,
)

# Tool / resource extraction patterns
_PY_TOOL_NAME = re.compile(
    r'@(?:mcp\.tool|app\.tool|server\.tool|tool)\s*\(\s*'
    r'(?:name\s*=\s*)?["\']([^"\']+)["\']'
)
_PY_TOOL_DEF = re.compile(
    r'@(?:mcp\.tool|app\.tool|server\.tool|tool)\s*\(.*?\)\s*\n\s*'
    r'(?:async\s+)?def\s+(\w+)',
    re.DOTALL,
)
_TS_TOOL_NAME = re.compile(
    r'\.tool\s*\(\s*["\']([^"\']+)["\']'
)
_PY_RESOURCE_NAME = re.compile(
    r'@(?:mcp\.resource|app\.resource|server\.resource|resource)\s*\(\s*'
    r'["\']([^"\']+)["\']'
)
_TS_RESOURCE_NAME = re.compile(
    r'\.resource\s*\(\s*["\']([^"\']+)["\']'
)

# Capability patterns
_CAPABILITY_PATTERNS: list[tuple[str, re.Pattern[str]]] = [
    ("tools", re.compile(r"\.tool\s*\(|@.*tool", re.IGNORECASE)),
    ("resources", re.compile(r"\.resource\s*\(|@.*resource", re.IGNORECASE)),
    ("prompts", re.compile(r"\.prompt\s*\(|@.*prompt", re.IGNORECASE)),
    ("logging", re.compile(r"logging\.getLogger|logger\s*=", re.IGNORECASE)),
    ("sampling", re.compile(r"sampling|create_message", re.IGNORECASE)),
]

_PY_EXTS = frozenset({".py", ".pyw"})
_TS_EXTS = frozenset({".ts", ".tsx", ".js", ".jsx", ".mjs"})
_SKIP_DIRS = frozenset({
    "node_modules", "__pycache__", ".git", ".venv", "venv",
})


def _detect_lang(filepath: Path) -> str:
    """Detect language from file extension."""
    ext = filepath.suffix.lower()
    if ext in _PY_EXTS:
        return "python"
    if ext in _TS_EXTS:
        return "typescript"
    return "unknown"


def _detect_transport(content: str) -> str:
    """Detect MCP transport type from source content."""
    if _STDIO_PATTERN.search(content):
        return "stdio"
    if _SSE_PATTERN.search(content):
        return "sse"
    if _HTTP_PATTERN.search(content):
        return "streamable-http"
    return "unknown"


def _detect_capabilities(content: str) -> list[str]:
    """Detect MCP capabilities from source code."""
    return [
        name for name, pat in _CAPABILITY_PATTERNS
        if pat.search(content)
    ]


def _extract_tools(content: str, language: str) -> list[str]:
    """Extract tool names from source code."""
    tools: list[str] = []
    if language == "python":
        for m in _PY_TOOL_NAME.finditer(content):
            tools.append(m.group(1))
        for m in _PY_TOOL_DEF.finditer(content):
            name = m.group(1)
            if name not in tools:
                tools.append(name)
    else:
        for m in _TS_TOOL_NAME.finditer(content):
            tools.append(m.group(1))
    return tools


def _extract_resources(content: str, language: str) -> list[str]:
    """Extract resource URIs from source code."""
    if language == "python":
        return [m.group(1) for m in _PY_RESOURCE_NAME.finditer(content)]
    return [m.group(1) for m in _TS_RESOURCE_NAME.finditer(content)]


def _extract_versions(content: str, framework: str) -> list[str]:
    """Extract version hints for the detected framework."""
    hints: list[str] = []
    for key, pat in _VERSION_PATTERNS.items():
        for m in pat.finditer(content):
            hints.append(m.group(1))
    return hints


def _fw_name_to_mcpfp(fw_name: str) -> str:
    """Normalize FrameworkInfo name to MCPFingerprint framework field."""
    mapping = {
        "FastMCP": "fastmcp",
        "mcp-python-sdk": "mcp-sdk",
        "@modelcontextprotocol/sdk": "mcp-sdk",
        "custom": "custom",
        "unknown": "unknown",
    }
    return mapping.get(fw_name, fw_name.lower())


def fingerprint_path(path: str | Path) -> MCPFingerprint:
    """Fingerprint an MCP server from a file or directory path.

    Args:
        path: Path to a source file or directory.

    Returns:
        MCPFingerprint with detected characteristics.
    """
    target = Path(path)
    fp = MCPFingerprint(file_path=str(target))

    if not target.exists():
        fp.errors.append(f"Path does not exist: {target}")
        return fp

    if target.is_dir():
        return _fingerprint_dir(target)

    return _fingerprint_single(target)


def _fingerprint_single(filepath: Path) -> MCPFingerprint:
    """Fingerprint a single source file."""
    fp = MCPFingerprint(file_path=str(filepath))

    try:
        content = filepath.read_text(encoding="utf-8", errors="replace")
    except (OSError, PermissionError) as exc:
        fp.errors.append(f"Cannot read file: {exc}")
        return fp

    fp.language = _detect_lang(filepath)

    # Use existing fingerprint_server for framework detection
    fw_info = fingerprint_server(content)
    fp.framework = _fw_name_to_mcpfp(fw_info.name)

    fp.transport = _detect_transport(content)
    fp.capabilities = _detect_capabilities(content)
    fp.version_hints = _extract_versions(content, fw_info.name)
    fp.tools = _extract_tools(content, fp.language)
    fp.resources = _extract_resources(content, fp.language)
    fp.tool_count = len(fp.tools)
    fp.resource_count = len(fp.resources)

    return fp


def _fingerprint_dir(directory: Path) -> MCPFingerprint:
    """Fingerprint a directory by aggregating all source files."""
    combined = MCPFingerprint(file_path=str(directory))
    all_tools: list[str] = []
    all_resources: list[str] = []
    all_caps: list[str] = []
    all_versions: list[str] = []

    for filepath in sorted(directory.rglob("*")):
        if not filepath.is_file():
            continue
        if filepath.suffix.lower() not in (_PY_EXTS | _TS_EXTS):
            continue
        if any(p in filepath.parts for p in _SKIP_DIRS):
            continue

        fp = _fingerprint_single(filepath)
        if fp.framework not in ("unknown", "custom") and combined.framework in ("unknown", "custom"):
            combined.framework = fp.framework
            combined.language = fp.language
        if fp.transport != "unknown" and combined.transport == "unknown":
            combined.transport = fp.transport
        all_tools.extend(fp.tools)
        all_resources.extend(fp.resources)
        all_caps.extend(fp.capabilities)
        all_versions.extend(fp.version_hints)

    combined.tools = list(dict.fromkeys(all_tools))
    combined.resources = list(dict.fromkeys(all_resources))
    combined.capabilities = list(dict.fromkeys(all_caps))
    combined.version_hints = list(dict.fromkeys(all_versions))
    combined.tool_count = len(combined.tools)
    combined.resource_count = len(combined.resources)

    if combined.framework == "unknown":
        combined.errors.append("No MCP server patterns detected")

    return combined


def format_mcpfingerprint_text(fp: MCPFingerprint) -> str:
    """Format MCPFingerprint as human-readable text."""
    lines: list[str] = []
    lines.append(f"\n{'=' * 50}")
    lines.append("  mcpaudit -- MCP Server Fingerprint")
    lines.append(f"{'=' * 50}")
    lines.append(f"  File:       {fp.file_path}")
    lines.append(f"  Framework:  {fp.framework}")
    lines.append(f"  Language:   {fp.language}")
    lines.append(f"  Transport:  {fp.transport}")
    lines.append(f"  Tools:      {fp.tool_count}")
    lines.append(f"  Resources:  {fp.resource_count}")

    if fp.version_hints:
        lines.append(f"  Versions:   {', '.join(fp.version_hints)}")
    if fp.capabilities:
        lines.append(f"  Capabilities: {', '.join(fp.capabilities)}")
    if fp.tools:
        lines.append("  Tool names:")
        for t in fp.tools:
            lines.append(f"    - {t}")
    if fp.resources:
        lines.append("  Resource URIs:")
        for r in fp.resources:
            lines.append(f"    - {r}")
    if fp.errors:
        lines.append("  Errors:")
        for e in fp.errors:
            lines.append(f"    ! {e}")

    lines.append(f"{'=' * 50}\n")
    return "\n".join(lines)


def format_mcpfingerprint_json(fp: MCPFingerprint) -> str:
    """Format MCPFingerprint as JSON."""
    import json as _json
    return _json.dumps(fp.to_dict(), indent=2, ensure_ascii=False)
