"""MCP server for the Exploit Intelligence Platform.

Exposes vulnerability and exploit intelligence to AI assistants via
the Model Context Protocol.  Supports both stdio and Streamable HTTP
transports (MCP spec 2025-03-26).

Usage:
    eip-mcp                                        # stdio (default)
    eip-mcp --transport streamable-http            # HTTP on 127.0.0.1:8080
    eip-mcp --transport streamable-http --port 9000
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import os
import sys
import time
from collections import deque
from threading import Lock
from typing import Any

import mcp.server.stdio
from mcp import types
from mcp.server.lowlevel import NotificationOptions, Server
from mcp.server.models import InitializationOptions

from eip_mcp import __version__
from eip_mcp import api_client
from eip_mcp import formatters
from eip_mcp import validators
from eip_mcp.validators import ValidationError

logging.basicConfig(level=logging.WARNING, stream=sys.stderr)
logger = logging.getLogger("eip-mcp")

_INSTRUCTIONS = """You have access to the Exploit Intelligence Platform (EIP) — 370K+ vulnerabilities and 105K+ exploits from NVD, CISA KEV, EPSS, ExploitDB, Metasploit, GitHub, and nomi-sec.

TOOL SELECTION GUIDE:
- "What vulns affect OpenSSH?" or any vendor/product query → search_vulnerabilities (has full-text query= param)
- Specific CVE ID mentioned → get_vulnerability (full brief with ranked exploits, AI analysis, MITRE ATT&CK)
- "Find all Metasploit modules" or filter exploits by source/language/author → search_exploits (structured filters ONLY — no free-text query)
- "Find reliable RCE exploits" → search_exploits with attack_type, complexity, reliability filters
- "Exploits for CVE-2024-XXXX" → PREFER get_vulnerability (includes exploits). Use search_exploits with cve= only if you need to page through 50+ exploits.
- Trojans or backdoored exploits → search_exploits with llm_classification=trojan
- All exploits from a vendor → search_exploits with vendor= filter
- Exploit authors → list_authors or get_author
- Vulnerability types or CWEs → list_cwes or get_cwe
- Which vendors have the most vulns → list_vendors
- "What products does Microsoft have?" or need exact product names → list_products
- EDB number or GHSA ID → lookup_alt_id to resolve to CVE
- Tech stack risk assessment → audit_stack
- Nuclei templates or Shodan/FOFA/Google dorks → get_nuclei_templates
- Pentest report finding → generate_finding
- Read exploit source code → get_exploit_code

IMPORTANT — search_vulnerabilities vs search_exploits:
- search_vulnerabilities supports free-text query= ("openssh", "log4j", "apache httpd"). START HERE for product research.
- search_exploits has NO query param. It only has structured filters (source, language, cve, vendor, product, attack_type, etc.).
- To find exploits for a product: search_vulnerabilities first → then get_vulnerability on interesting CVEs (exploits are included).

KEY CONCEPTS:
- EPSS: Probability (0-100%) that a CVE will be exploited in the next 30 days. >90% is very high risk.
- CVSS: Severity score (0-10). Critical >= 9.0, High >= 7.0, Medium >= 4.0.
- KEV: CISA Known Exploited Vulnerabilities — confirmed actively exploited in the wild.
- Exploit ranking: Metasploit modules > verified ExploitDB > GitHub by stars. Trojans flagged at bottom.
- AI Analysis: 61K+ exploits analyzed. Results include attack_type (RCE, SQLi, XSS, DoS, LPE, auth_bypass, info_leak), complexity (trivial/simple/moderate/complex), reliability (reliable/unreliable/untested), target_software, MITRE ATT&CK techniques, and deception_indicators for trojans.
MULTI-TOOL WORKFLOWS:
- "Tell me about CVE-XXXX": get_vulnerability → full brief including exploits with MITRE techniques
- "What vulns affect Apache?": search_vulnerabilities(query="apache httpd", has_exploits=true) → get_vulnerability on top hits
- "What should I worry about?": audit_stack with technologies → get_vulnerability on top findings
- "Is there a Metasploit module for X?": get_vulnerability → check METASPLOIT MODULES section
- "Write a finding for my report": generate_finding with CVE, target, and notes
"""

server = Server("eip-mcp", version=__version__, instructions=_INSTRUCTIONS)


# ---------------------------------------------------------------------------
# Local safety controls
# ---------------------------------------------------------------------------


class LocalRateLimitError(Exception):
    """Raised when local (process-level) rate limiting is exceeded."""

    def __init__(self, retry_after_s: float):
        self.retry_after_s = max(0.0, float(retry_after_s))
        super().__init__(f"Rate limited. Retry after {self.retry_after_s:.1f}s")


class RateLimiter:
    """Simple sliding-window limiter for tool calls.

    This limits *tool invocations* (not HTTP requests) per process.
    """

    def __init__(self, max_calls: int, period_s: float):
        self.max_calls = int(max_calls)
        self.period_s = float(period_s)
        self._times: deque[float] = deque()
        # Use a thread lock to avoid binding this limiter to a specific event loop.
        self._lock = Lock()

    def enabled(self) -> bool:
        return self.max_calls > 0 and self.period_s > 0

    def check_or_raise(self, cost: int = 1) -> None:
        if not self.enabled():
            return
        now = time.monotonic()
        with self._lock:
            cutoff = now - self.period_s
            while self._times and self._times[0] < cutoff:
                self._times.popleft()
            if len(self._times) + cost > self.max_calls:
                retry_after = (self._times[0] + self.period_s) - now
                raise LocalRateLimitError(retry_after)
            for _ in range(cost):
                self._times.append(now)


def _int_env(name: str, default: int) -> int:
    raw = os.environ.get(name)
    if raw is None or raw == "":
        return default
    try:
        return int(raw)
    except ValueError:
        return default


def _float_env(name: str, default: float) -> float:
    raw = os.environ.get(name)
    if raw is None or raw == "":
        return default
    try:
        return float(raw)
    except ValueError:
        return default


_MAX_CONCURRENCY = max(0, _int_env("EIP_MCP_MAX_CONCURRENCY", 4))
_CONCURRENCY_SEM: asyncio.Semaphore | None = None

_MAX_CALLS = max(0, _int_env("EIP_MCP_MAX_CALLS", 60))
_CALL_PERIOD_S = max(0.0, _float_env("EIP_MCP_CALL_PERIOD_SECONDS", 60.0))
_RATE_LIMITER = RateLimiter(_MAX_CALLS, _CALL_PERIOD_S)

# ---------------------------------------------------------------------------
# Tool definitions
# ---------------------------------------------------------------------------

TOOLS = [
    types.Tool(
        name="search_vulnerabilities",
        description=(
            "Search the Exploit Intelligence Platform for vulnerabilities (CVEs). "
            "Returns a list of matching CVEs with CVSS scores, EPSS exploitation "
            "probability, exploit counts, and CISA KEV status. Supports full-text "
            "search, severity/vendor/product/ecosystem/CWE filters, and CVSS/EPSS "
            "thresholds. Use this for broad searches like 'apache vulnerabilities' "
            "or filtered queries like 'critical Fortinet CVEs with exploits'. "
            "Examples: query='apache httpd' with has_exploits=true; "
            "vendor='fortinet' with severity='critical' and is_kev=true sorted by epss_desc; "
            "cwe='89' with min_cvss=9 for critical SQL injection CVEs."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "Search keywords (e.g. 'apache httpd', 'log4j'). Optional if filters are provided."},
                "severity": {"type": "string", "enum": ["critical", "high", "medium", "low"], "description": "Filter by severity level"},
                "has_exploits": {"type": "boolean", "description": "Only return CVEs with public exploit code"},
                "is_kev": {"type": "boolean", "description": "Only return CISA Known Exploited Vulnerabilities"},
                "has_nuclei": {"type": "boolean", "description": "Only return CVEs with Nuclei scanner templates"},
                "vendor": {"type": "string", "description": "Filter by vendor name (e.g. 'microsoft', 'fortinet')"},
                "product": {"type": "string", "description": "Filter by product name (e.g. 'exchange', 'pan-os')"},
                "ecosystem": {"type": "string", "enum": ["npm", "pip", "maven", "go", "crates", "nuget", "rubygems", "composer"], "description": "Filter by package ecosystem"},
                "cwe": {"type": "string", "description": "Filter by CWE ID (e.g. '79' or 'CWE-79')"},
                "min_cvss": {"type": "number", "description": "Minimum CVSS score (0-10)"},
                "min_epss": {"type": "number", "description": "Minimum EPSS score (0-1)"},
                "year": {"type": "integer", "description": "Filter by CVE year (e.g. 2024)"},
                "date_from": {"type": "string", "description": "Start date for CVE publication (YYYY-MM-DD)"},
                "date_to": {"type": "string", "description": "End date for CVE publication (YYYY-MM-DD)"},
                "sort": {"type": "string", "enum": ["newest", "oldest", "cvss_desc", "epss_desc", "relevance"], "description": "Sort order (default: newest)"},
                "page": {"type": "integer", "description": "Page number (default: 1)"},
                "per_page": {"type": "integer", "description": "Results per page (1-25, default: 10)"},
            },
        },
    ),
    types.Tool(
        name="get_vulnerability",
        description=(
            "Get a full intelligence brief for a specific vulnerability. Accepts both "
            "CVE-IDs (e.g. CVE-2024-3400) and EIP-IDs (e.g. EIP-2026-12345 for pre-CVE entries). "
            "Returns detailed information including CVSS score and vector, EPSS exploitation "
            "probability, CISA KEV status, description, affected products, ranked exploits "
            "(grouped by Metasploit modules, verified ExploitDB, GitHub PoCs, and trojans), "
            "Nuclei scanner templates with recon dorks, alternate identifiers, and references. "
            "Exploits are ranked by quality: Metasploit modules first (peer-reviewed), "
            "then verified ExploitDB, then GitHub by stars. Trojans are flagged at the bottom."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "cve_id": {"type": "string", "description": "CVE or EIP identifier (e.g. 'CVE-2024-3400' or 'EIP-2026-12345')"},
            },
            "required": ["cve_id"],
        },
    ),
    types.Tool(
        name="get_exploit_code",
        description=(
            "Retrieve the source code of a specific exploit by its platform ID. "
            "IMPORTANT: Use the platform's internal ID shown as [id=XXXXX] in results, "
            "NOT the ExploitDB number (EDB-XXXXX). These are different numbering systems. "
            "Returns code from the exploit archive. If no file_path is specified, "
            "auto-selects the most relevant code file. Use this to analyze exploit "
            "mechanics, understand attack techniques, or review PoC code."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "exploit_id": {"type": "integer", "description": "Platform exploit ID (the [id=XXXXX] number from results — NOT the EDB number)"},
                "file_path": {"type": "string", "description": "Specific file path to view (optional — auto-selects if omitted)"},
            },
            "required": ["exploit_id"],
        },
    ),
    types.Tool(
        name="get_nuclei_templates",
        description=(
            "Get Nuclei scanner templates and recon dorks for a vulnerability. Returns "
            "template metadata, severity, verification status, tags, and ready-to-use "
            "Shodan, FOFA, and Google dork queries for target identification. Accepts "
            "both CVE-IDs and EIP-IDs. Use this to plan scanning or reconnaissance."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "cve_id": {"type": "string", "description": "CVE or EIP identifier (e.g. 'CVE-2024-27198')"},
            },
            "required": ["cve_id"],
        },
    ),
    types.Tool(
        name="get_platform_stats",
        description=(
            "Get platform-wide statistics from the Exploit Intelligence Platform. "
            "Returns total counts of vulnerabilities, exploits, KEV entries, Nuclei "
            "templates, vendors, and authors, plus the last data update timestamp."
        ),
        inputSchema={"type": "object", "properties": {}},
    ),
    types.Tool(
        name="check_health",
        description=(
            "Check the EIP API health and data freshness. Returns database status "
            "and timestamps for each of the 11 ingestion sources (NVD, KEV, EPSS, "
            "ExploitDB, GitHub, Metasploit, etc.)."
        ),
        inputSchema={"type": "object", "properties": {}},
    ),
    types.Tool(
        name="search_exploits",
        description=(
            "Browse and filter exploits using STRUCTURED FILTERS ONLY (no free-text query). "
            "Use this to filter by source (github, metasploit, exploitdb, nomisec), language "
            "(python, ruby, etc.), LLM classification (working_poc, trojan, suspicious, scanner, "
            "stub, writeup), author, min stars, code availability, CVE ID, vendor, or product. "
            "Also filter by AI analysis: attack_type (RCE, SQLi, XSS, DoS, LPE, auth_bypass, "
            "info_leak), complexity (trivial/simple/moderate/complex), reliability "
            "(reliable/unreliable/untested), requires_auth. "
            "NOTE: To search by product name (e.g. 'OpenSSH', 'Apache'), use "
            "search_vulnerabilities instead — it has free-text query and get_vulnerability "
            "already includes exploits in the response. "
            "Examples: source='metasploit' for all Metasploit modules; "
            "attack_type='RCE' with reliability='reliable' for weaponizable RCE exploits; "
            "cve='CVE-2024-3400' for all exploits targeting a specific CVE; "
            "vendor='mitel' for all Mitel exploits."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "source": {"type": "string", "description": "Filter by source: github (direct GitHub crawl, ~400), metasploit (~3,900 modules), exploitdb (~52K), nomisec (GitHub PoCs via nomi-sec, ~19K — most GitHub exploits are here), writeup (reference writeups)"},
                "language": {"type": "string", "description": "Filter by language: python, ruby, go, c, etc."},
                "llm_classification": {"type": "string", "description": "Filter by LLM classification: working_poc, trojan, suspicious, scanner, stub, writeup, tool"},
                "attack_type": {"type": "string", "description": "Filter by attack type from AI analysis: RCE, SQLi, XSS, DoS, LPE, auth_bypass, info_leak, deserialization, other"},
                "complexity": {"type": "string", "description": "Filter by exploit complexity: trivial, simple, moderate, complex"},
                "reliability": {"type": "string", "description": "Filter by exploit reliability: reliable, unreliable, untested"},
                "requires_auth": {"type": "boolean", "description": "Filter by whether exploit requires authentication"},
                "author": {"type": "string", "description": "Filter by author name"},
                "min_stars": {"type": "integer", "description": "Minimum GitHub stars"},
                "has_code": {"type": "boolean", "description": "Only exploits with downloadable code"},
                "cve": {"type": "string", "description": "Filter by CVE ID (e.g. 'CVE-2024-3400') — returns all exploits for that CVE"},
                "vendor": {"type": "string", "description": "Filter by vendor name (e.g. 'mitel', 'fortinet') — returns exploits for all CVEs affecting that vendor"},
                "product": {"type": "string", "description": "Filter by product name (e.g. 'micollab', 'pan-os')"},
                "sort": {"type": "string", "enum": ["newest", "stars_desc"], "description": "Sort order"},
                "page": {"type": "integer", "description": "Page number (default: 1)"},
                "per_page": {"type": "integer", "description": "Results per page (1-25, default: 10)"},
            },
        },
    ),
    types.Tool(
        name="list_authors",
        description=(
            "List exploit authors/researchers ranked by exploit count. Returns the top "
            "security researchers with their exploit counts and handles. "
            "Use this when asked 'who are the top exploit authors?' or 'who writes the most exploits?'"
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "page": {"type": "integer", "description": "Page number (default: 1)"},
                "per_page": {"type": "integer", "description": "Results per page (1-50, default: 25)"},
            },
        },
    ),
    types.Tool(
        name="get_author",
        description=(
            "Get an exploit author's profile with all their exploits. Returns author name, "
            "handle, total exploit count, activity start date, and a paginated list of their "
            "exploits with CVE context. Use this when asked about a specific researcher "
            "like 'show me all exploits by Chocapikk'."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "author_name": {"type": "string", "description": "Author name (e.g. 'Chocapikk')"},
            },
            "required": ["author_name"],
        },
    ),
    types.Tool(
        name="list_cwes",
        description=(
            "List CWE (Common Weakness Enumeration) categories ranked by vulnerability count. "
            "Returns CWE IDs, names, short labels, exploit likelihood, and how many CVEs "
            "have that weakness. Use this when asked 'what are the most common vulnerability types?'"
        ),
        inputSchema={"type": "object", "properties": {}},
    ),
    types.Tool(
        name="get_cwe",
        description=(
            "Get details for a specific CWE including full name, description, exploit "
            "likelihood, parent CWE, and total vulnerability count. "
            "Example: cwe_id='CWE-79' returns details about Cross-Site Scripting."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "cwe_id": {"type": "string", "description": "CWE identifier (e.g. 'CWE-79' or '79')"},
            },
            "required": ["cwe_id"],
        },
    ),
    types.Tool(
        name="list_vendors",
        description=(
            "List software vendors ranked by vulnerability count. Returns the top 200 "
            "vendors with their total CVE counts. Use this when asked 'which vendors have "
            "the most vulnerabilities?' or to understand the threat landscape by vendor."
        ),
        inputSchema={"type": "object", "properties": {}},
    ),
    types.Tool(
        name="list_products",
        description=(
            "List products for a specific vendor with vulnerability counts. Use this "
            "to discover exact product names for filtering. Product names in the "
            "database use CPE conventions (e.g. 'exchange_server' not 'exchange', "
            "'windows_10' not 'windows 10'). Example: vendor='microsoft' returns "
            "products like exchange_server, windows_10, office, edge_chromium."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "vendor": {"type": "string", "description": "Vendor name (e.g. 'microsoft', 'apache', 'fortinet')"},
            },
            "required": ["vendor"],
        },
    ),
    types.Tool(
        name="lookup_alt_id",
        description=(
            "Look up a vulnerability by an alternate identifier such as an ExploitDB "
            "ID (EDB-XXXXX) or GitHub Security Advisory ID (GHSA-XXXXX). Returns the "
            "matching CVE-ID with basic severity info. Use this when you have an EDB "
            "number or GHSA ID and need to find the corresponding CVE."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "alt_id": {"type": "string", "description": "Alternate ID (e.g. 'EDB-48537', 'GHSA-jfh8-c2jp-5v3q')"},
            },
            "required": ["alt_id"],
        },
    ),
    types.Tool(
        name="audit_stack",
        description=(
            "Audit a technology stack for exploitable vulnerabilities. Accepts a "
            "comma-separated list of technologies (max 5) and searches for critical/ "
            "high severity CVEs with public exploits for each one, sorted by EPSS "
            "exploitation probability. Use this when a user describes their "
            "infrastructure and wants to know what to patch first. "
            "Example: technologies='nginx, postgresql, node.js' returns a risk-sorted "
            "list of exploitable CVEs grouped by technology."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "technologies": {
                    "type": "string",
                    "description": "Comma-separated list of technologies (e.g. 'nginx, postgresql, node.js'). Max 5.",
                },
            },
            "required": ["technologies"],
        },
    ),
    types.Tool(
        name="generate_finding",
        description=(
            "Generate a pentest report finding in Markdown format for a specific "
            "vulnerability. Fetches full detail and formats it as a professional finding "
            "with severity, CVSS, description, affected products, exploit availability, "
            "and references. Accepts both CVE-IDs and EIP-IDs. Optionally include the "
            "target system tested and tester notes. The output is ready to paste into "
            "a pentest report. "
            "Example: cve_id='CVE-2024-3400', target='fw.corp.example.com', "
            "notes='Confirmed RCE via GlobalProtect gateway'."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "cve_id": {"type": "string", "description": "CVE or EIP identifier (e.g. 'CVE-2024-3400')"},
                "target": {"type": "string", "description": "Target system tested (e.g. 'fw.corp.example.com'). Optional."},
                "notes": {"type": "string", "description": "Tester notes to include in the finding. Optional."},
            },
            "required": ["cve_id"],
        },
    ),
]


@server.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    return TOOLS


@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any] | None):
    """Route tool calls to handlers with validation and error handling.

    Notes:
    - Tool calls run with local concurrency and rate limiting.
    - Sync work (HTTP calls) executes in a worker thread to avoid blocking the event loop.
    """
    args = arguments or {}

    # Composite tools make multiple backend API calls per invocation.
    # Charge them proportionally so one client can't amplify load.
    _TOOL_COST = {"audit_stack": 10}  # up to 5 techs × 2 API calls

    try:
        _RATE_LIMITER.check_or_raise(cost=_TOOL_COST.get(name, 1))

        global _CONCURRENCY_SEM
        if _MAX_CONCURRENCY and _CONCURRENCY_SEM is None:
            _CONCURRENCY_SEM = asyncio.Semaphore(_MAX_CONCURRENCY)

        if _CONCURRENCY_SEM is None:
            result = await asyncio.to_thread(_dispatch, name, args)
        else:
            async with _CONCURRENCY_SEM:
                result = await asyncio.to_thread(_dispatch, name, args)

        return _tool_result(result, is_error=False)

    except LocalRateLimitError as exc:
        return _tool_result(
            f"Rate limited locally to protect reliability. Retry after {exc.retry_after_s:.1f}s.",
            is_error=True,
        )
    except ValidationError as exc:
        return _tool_result(f"Validation error: {exc}", is_error=True)
    except api_client.APIError as exc:
        if exc.status_code == 404:
            if "cwe_id" in args:
                cwe_val = args["cwe_id"]
                cwe_display = cwe_val if str(cwe_val).upper().startswith("CWE-") else f"CWE-{cwe_val}"
                return _tool_result(f"{cwe_display} not found in database.", is_error=True)
            if "author_name" in args:
                return _tool_result(f"Author '{args['author_name']}' not found.", is_error=True)
            if "alt_id" in args:
                return _tool_result(f"No CVE found for '{args['alt_id']}'.", is_error=True)
            ident = args.get("cve_id", args.get("exploit_id", "?"))
            return _tool_result(f"Not found: {ident}", is_error=True)
        if exc.status_code == 429:
            return _tool_result(f"Rate limited by the EIP API. {exc.message}", is_error=True)
        if exc.status_code == 0:
            return _tool_result(f"Network error contacting the EIP API. {exc.message}", is_error=True)
        return _tool_result(f"API error: {exc.message}", is_error=True)
    except Exception:
        logger.exception("Unexpected error in tool %s", name)
        return _tool_result("Internal error processing request. Please try again.", is_error=True)


def _tool_result(text: str, *, is_error: bool) -> Any:
    """Build a CallTool result with explicit error signaling when supported."""
    content = [types.TextContent(type="text", text=text)]

    # Newer MCP SDKs expose CallToolResult (with an isError field) which clients
    # use to distinguish normal output from failures.
    if hasattr(types, "CallToolResult"):
        try:
            return types.CallToolResult(content=content, isError=is_error)
        except TypeError:
            # Some SDK versions may use snake_case; keep compatibility.
            return types.CallToolResult(content=content, is_error=is_error)  # type: ignore[arg-type]

    # Fallback for older SDKs.
    return content


# ---------------------------------------------------------------------------
# Tool dispatch
# ---------------------------------------------------------------------------

def _dispatch(name: str, args: dict[str, Any]) -> str:
    if name == "search_vulnerabilities":
        return _tool_search(args)
    elif name == "get_vulnerability":
        return _tool_get_vuln(args)
    elif name == "get_exploit_code":
        return _tool_get_code(args)
    elif name == "get_nuclei_templates":
        return _tool_get_nuclei(args)
    elif name == "get_platform_stats":
        return _tool_stats()
    elif name == "check_health":
        return _tool_health()
    elif name == "search_exploits":
        return _tool_search_exploits(args)
    elif name == "list_authors":
        return _tool_list_authors(args)
    elif name == "get_author":
        return _tool_get_author(args)
    elif name == "list_cwes":
        return _tool_list_cwes()
    elif name == "get_cwe":
        return _tool_get_cwe(args)
    elif name == "list_vendors":
        return _tool_list_vendors()
    elif name == "list_products":
        return _tool_list_products(args)
    elif name == "lookup_alt_id":
        return _tool_lookup_alt_id(args)
    elif name == "audit_stack":
        return _tool_audit_stack(args)
    elif name == "generate_finding":
        return _tool_generate_finding(args)
    else:
        raise ValidationError(f"Unknown tool: {name}")


# ---------------------------------------------------------------------------
# Core tool handlers
# ---------------------------------------------------------------------------

def _tool_search(args: dict) -> str:
    params: dict[str, Any] = {}
    if "query" in args and args["query"]:
        params["q"] = validators.validate_query(args["query"])
    if "severity" in args and args["severity"]:
        params["severity"] = validators.validate_severity(args["severity"])
    if args.get("has_exploits"):
        params["has_exploits"] = True
    if args.get("is_kev"):
        params["is_kev"] = True
    if args.get("has_nuclei"):
        params["has_nuclei"] = True
    if "vendor" in args and args["vendor"]:
        params["vendor"] = validators.validate_vendor(args["vendor"])
    if "product" in args and args["product"]:
        params["product"] = validators.validate_product(args["product"])
    if "ecosystem" in args and args["ecosystem"]:
        params["ecosystem"] = validators.validate_ecosystem(args["ecosystem"])
    if "cwe" in args and args["cwe"]:
        params["cwe"] = validators.validate_cwe(args["cwe"])
    if "min_cvss" in args and args["min_cvss"] is not None:
        params["min_cvss"] = validators.validate_cvss(args["min_cvss"])
    if "min_epss" in args and args["min_epss"] is not None:
        params["min_epss"] = validators.validate_epss(args["min_epss"])
    if "year" in args and args["year"] is not None:
        params["year"] = validators.validate_year(args["year"])
    if "date_from" in args and args["date_from"]:
        params["date_from"] = validators.validate_date(args["date_from"])
    if "date_to" in args and args["date_to"]:
        params["date_to"] = validators.validate_date(args["date_to"])
    if "sort" in args and args["sort"]:
        params["sort"] = validators.validate_sort(args["sort"])
    if "page" in args and args["page"] is not None:
        params["page"] = validators.validate_page(args["page"])
    if "per_page" in args and args["per_page"] is not None:
        params["per_page"] = validators.validate_per_page(args["per_page"])
    else:
        params["per_page"] = 10

    if not params or all(k in ("per_page", "page") for k in params):
        raise ValidationError("Provide a search query or at least one filter.")

    data = api_client.search_vulns(params)
    return formatters.format_search_results(data)


def _tool_get_vuln(args: dict) -> str:
    vuln_id = validators.validate_vuln_id(args.get("cve_id", ""))
    data = api_client.get_vuln_detail(vuln_id)
    return formatters.format_vuln_detail(data)


def _tool_get_code(args: dict) -> str:
    exploit_id = validators.validate_exploit_id(args.get("exploit_id", ""))
    file_path = args.get("file_path")

    if file_path:
        file_path = validators.validate_file_path(file_path)
    else:
        # Auto-select: list files and pick the best one
        files = api_client.list_exploit_files(exploit_id)
        if not files:
            return f"No code files found for exploit {exploit_id}."
        file_path = _pick_main_file(files)
        if not file_path:
            return formatters.format_exploit_files(files, exploit_id) + "\n\nSpecify file_path to view a specific file."

    code = api_client.get_exploit_code(exploit_id, file_path)
    return formatters.format_exploit_code(code, file_path)


def _tool_get_nuclei(args: dict) -> str:
    cve_id = validators.validate_vuln_id(args.get("cve_id", ""))
    data = api_client.get_vuln_detail(cve_id)
    templates = data.get("nuclei_templates") or []
    if not templates:
        return f"No Nuclei templates found for {cve_id}."
    return formatters.format_nuclei_templates(templates, header=True)


def _tool_stats() -> str:
    data = api_client.get_stats()
    return formatters.format_stats(data)


def _tool_health() -> str:
    data = api_client.get_health()
    return formatters.format_health(data)


# ---------------------------------------------------------------------------
# Browse tool handlers (new endpoints)
# ---------------------------------------------------------------------------

def _tool_search_exploits(args: dict) -> str:
    params: dict[str, Any] = {}
    if args.get("source"):
        params["source"] = validators.validate_source(args["source"])
    if args.get("language"):
        params["language"] = validators.validate_language(args["language"])
    if args.get("llm_classification"):
        params["llm_classification"] = validators.validate_llm_classification(args["llm_classification"])
    if args.get("attack_type"):
        params["attack_type"] = validators.validate_attack_type(args["attack_type"])
    if args.get("complexity"):
        params["complexity"] = validators.validate_complexity(args["complexity"])
    if args.get("reliability"):
        params["reliability"] = validators.validate_reliability(args["reliability"])
    if "requires_auth" in args and args["requires_auth"] is not None:
        if not isinstance(args["requires_auth"], bool):
            raise ValidationError("requires_auth must be true or false")
        params["requires_auth"] = args["requires_auth"]
    if args.get("author"):
        params["author"] = validators.validate_query(args["author"])
    if args.get("min_stars") is not None:
        params["min_stars"] = validators.validate_min_stars(args["min_stars"])
    if "has_code" in args and args["has_code"] is not None:
        if not isinstance(args["has_code"], bool):
            raise ValidationError("has_code must be true or false")
        params["has_code"] = args["has_code"]
    if args.get("cve"):
        params["cve"] = validators.validate_cve_id(args["cve"])
    if args.get("vendor"):
        params["vendor"] = validators.validate_vendor(args["vendor"])
    if args.get("product"):
        params["product"] = validators.validate_product(args["product"])
    if args.get("sort"):
        s = args["sort"]
        if s not in ("newest", "stars_desc"):
            raise ValidationError("sort must be 'newest' or 'stars_desc'")
        params["sort"] = s
    if "page" in args and args["page"] is not None:
        params["page"] = validators.validate_page(args["page"])
    if args.get("per_page") is not None:
        params["per_page"] = validators.validate_per_page(args["per_page"])
    else:
        params["per_page"] = 10

    data = api_client.search_exploits(params)
    return formatters.format_exploit_search(data)


def _tool_list_authors(args: dict) -> str:
    params: dict[str, Any] = {}
    if "page" in args and args["page"] is not None:
        params["page"] = validators.validate_page(args["page"])
    if args.get("per_page") is not None:
        try:
            val = int(args["per_page"])
        except (TypeError, ValueError):
            raise ValidationError(f"per_page must be an integer, got: '{args['per_page']}'")
        if val < 1 or val > 50:
            raise ValidationError(f"per_page must be 1-50, got: {val}")
        params["per_page"] = val
    else:
        params["per_page"] = 25
    data = api_client.list_authors(params)
    return formatters.format_authors_list(data)


def _tool_get_author(args: dict) -> str:
    name = validators.validate_query(args.get("author_name", ""))
    if not name:
        raise ValidationError("author_name is required")
    data = api_client.get_author(name)
    return formatters.format_author_detail(data)


def _tool_list_cwes() -> str:
    data = api_client.list_cwes()
    return formatters.format_cwe_list(data)


def _tool_get_cwe(args: dict) -> str:
    cwe_id = validators.validate_cwe(args.get("cwe_id", ""))
    data = api_client.get_cwe(f"CWE-{cwe_id}")
    return formatters.format_cwe_detail(data)


def _tool_list_vendors() -> str:
    data = api_client.list_vendors()
    return formatters.format_vendors_list(data)


def _tool_list_products(args: dict) -> str:
    vendor = validators.validate_vendor(args.get("vendor", ""))
    if not vendor:
        raise ValidationError("vendor is required")
    data = api_client.list_vendor_products(vendor)
    return formatters.format_products_list(data)


def _tool_lookup_alt_id(args: dict) -> str:
    alt_id = validators.clean_string(args.get("alt_id", ""), max_len=100)
    if not alt_id:
        raise ValidationError("alt_id is required")
    data = api_client.lookup_alt_id(alt_id)
    cve = data.get("cve_id") or data.get("eip_id", "?")
    title = data.get("title") or "No title"
    sev = (data.get("severity_label") or "?").upper()
    cvss = data.get("cvss_v3_score")
    cvss_str = f"CVSS: {cvss:.1f}" if cvss is not None else ""
    return f"{data.get('alt_id', alt_id)} → {cve}\n{title}  [{sev}]  {cvss_str}"


# ---------------------------------------------------------------------------
# Smart composite tools
# ---------------------------------------------------------------------------

def _tool_audit_stack(args: dict) -> str:
    techs = validators.validate_technologies(args.get("technologies", ""))
    results: dict[str, dict] = {}
    base_params = {
        "has_exploits": True,
        "min_cvss": 7.0,
        "sort": "epss_desc",
        "per_page": 10,
    }
    for tech in techs:
        try:
            # Strategy 1: exact product match (works for single-word CPE names)
            data = api_client.search_vulns({**base_params, "product": tech})
            if data.get("total", 0) > 0:
                results[tech] = data
                continue
            # Strategy 2: full-text search (catches natural names like "apache httpd")
            data = api_client.search_vulns({**base_params, "q": tech})
            results[tech] = data
        except api_client.APIError:
            results[tech] = {"total": 0, "items": []}
    return formatters.format_stack_audit(results)


def _tool_generate_finding(args: dict) -> str:
    cve_id = validators.validate_vuln_id(args.get("cve_id", ""))
    target = args.get("target")
    notes = args.get("notes")
    if target:
        target = validators.validate_query(target)
    if notes:
        notes = validators.validate_query(notes)
    data = api_client.get_vuln_detail(cve_id)
    return formatters.format_finding(data, target, notes)


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

def _pick_main_file(files: list[dict]) -> str | None:
    """Auto-select the most relevant exploit file."""
    code_exts = {".py", ".rb", ".go", ".c", ".cpp", ".java", ".js", ".pl", ".php", ".sh", ".ps1"}
    code_files = [f for f in files if any(f["path"].endswith(ext) for ext in code_exts)]

    for pattern in ("exploit", "poc", "main", "scan", "vuln", "rce"):
        for f in code_files:
            if pattern in f["path"].lower():
                return f["path"]

    if code_files:
        return max(code_files, key=lambda f: f.get("size", 0))["path"]

    non_readme = [f for f in files if "readme" not in f["path"].lower()]
    if non_readme:
        return max(non_readme, key=lambda f: f.get("size", 0))["path"]

    return files[0]["path"] if files else None


# ---------------------------------------------------------------------------
# Entry point — stdio transport
# ---------------------------------------------------------------------------

async def _run_stdio() -> None:
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="eip-mcp",
                server_version=__version__,
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )


# ---------------------------------------------------------------------------
# Entry point — Streamable HTTP transport (MCP spec 2025-03-26)
# ---------------------------------------------------------------------------

def _run_http_server(
    *,
    host: str,
    port: int,
    stateless: bool,
) -> None:
    """Start the MCP server with Streamable HTTP transport.

    All HTTP-specific imports are deferred so that ``pip install eip-mcp``
    (without the ``[http]`` extra) keeps working for stdio-only users.
    """
    try:
        import contextlib
        from collections.abc import AsyncIterator

        import uvicorn
        from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
        from mcp.server.transport_security import TransportSecuritySettings
        from starlette.applications import Starlette
        from starlette.middleware.cors import CORSMiddleware
        from starlette.responses import JSONResponse, Response
        from starlette.routing import Mount
        from starlette.types import Receive, Scope, Send
    except ImportError as exc:
        logger.error(
            "Streamable HTTP transport requires additional packages.  "
            "Install them with:  pip install 'eip-mcp[http]'"
        )
        raise SystemExit(1) from exc

    # ------------------------------------------------------------------
    # DNS-rebinding protection  (MCP spec §Streamable HTTP / Security)
    #
    # The Host header MUST match one of the expected values.  The Origin
    # header (sent by browsers) MUST be in the allow-list when present;
    # native MCP clients omit it, which is accepted by default.
    # ------------------------------------------------------------------
    allowed_hosts: set[str] = set()
    allowed_hosts.add(f"{host}:{port}")
    if host in ("127.0.0.1", "localhost", "::1", "0.0.0.0", "::"):
        for alias in ("127.0.0.1", "localhost"):
            allowed_hosts.add(f"{alias}:{port}")

    extra_hosts = os.environ.get("EIP_MCP_ALLOWED_HOSTS", "")
    for h in extra_hosts.split(","):
        h = h.strip()
        if h:
            allowed_hosts.add(h)

    allowed_origins: list[str] = []
    extra_origins = os.environ.get("EIP_MCP_ALLOWED_ORIGINS", "")
    for o in extra_origins.split(","):
        o = o.strip()
        if o:
            allowed_origins.append(o)

    security_settings = TransportSecuritySettings(
        enable_dns_rebinding_protection=True,
        allowed_hosts=sorted(allowed_hosts),
        allowed_origins=allowed_origins,
    )

    # ------------------------------------------------------------------
    # Session manager
    # ------------------------------------------------------------------
    sm_kwargs: dict[str, Any] = dict(
        app=server,
        json_response=False,
        stateless=stateless,
        security_settings=security_settings,
    )
    # session_idle_timeout was added in a later SDK release; use it when
    # available so stateful sessions auto-expire after 30 minutes of idle.
    import inspect
    _sm_params = inspect.signature(StreamableHTTPSessionManager.__init__).parameters
    if "session_idle_timeout" in _sm_params and not stateless:
        sm_kwargs["session_idle_timeout"] = 1800.0

    session_manager = StreamableHTTPSessionManager(**sm_kwargs)

    async def handle_mcp(scope: Scope, receive: Receive, send: Send) -> None:
        await session_manager.handle_request(scope, receive, send)

    @contextlib.asynccontextmanager
    async def lifespan(_app: Starlette) -> AsyncIterator[None]:
        async with session_manager.run():
            logger.warning(
                "eip-mcp v%s  transport=streamable-http  "
                "endpoint=http://%s:%d/mcp  mode=%s",
                __version__,
                host,
                port,
                "stateless" if stateless else "stateful",
            )
            yield

    # ------------------------------------------------------------------
    # Request body size limit (1 MB)
    # ------------------------------------------------------------------
    _MAX_BODY_BYTES = 1_048_576

    _BODY_TOO_LARGE = JSONResponse(
        {"jsonrpc": "2.0", "id": "server-error",
         "error": {"code": -32600, "message": "Request too large"}},
        status_code=413,
    )

    class _BodySizeLimitMiddleware:
        """Reject HTTP requests whose Content-Length exceeds the limit."""

        def __init__(self, app: Any) -> None:
            self.app = app

        async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
            if scope["type"] == "http":
                for key, value in scope.get("headers", []):
                    if key == b"content-length":
                        try:
                            if int(value) > _MAX_BODY_BYTES:
                                await _BODY_TOO_LARGE(scope, receive, send)
                                return
                        except (ValueError, TypeError):
                            pass
                        break
            await self.app(scope, receive, send)

    # ------------------------------------------------------------------
    # Suppress Pydantic/MCP-SDK validation internals in error responses
    # ------------------------------------------------------------------
    _GENERIC_PARSE_ERROR = JSONResponse(
        {"jsonrpc": "2.0", "id": "server-error",
         "error": {"code": -32700, "message": "Invalid request"}},
        status_code=400,
    )

    class _ErrorSanitiserMiddleware:
        """Replace verbose 400/422 error bodies with a generic JSON-RPC error.

        The MCP SDK returns detailed Pydantic validation errors for
        malformed JSON-RPC messages.  These leak internal class names and
        framework versions — replace them with a generic message.
        """

        def __init__(self, app: Any) -> None:
            self.app = app

        async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
            if scope["type"] != "http":
                await self.app(scope, receive, send)
                return

            response_started = False
            status_code = 200

            async def send_wrapper(message: Any) -> None:
                nonlocal response_started, status_code
                if message["type"] == "http.response.start":
                    status_code = message.get("status", 200)
                    if status_code in (400, 422):
                        # Suppress — we'll send our own response instead.
                        response_started = True
                        return
                if message["type"] == "http.response.body" and response_started:
                    await _GENERIC_PARSE_ERROR(scope, receive, send)
                    return
                await send(message)

            await self.app(scope, receive, send_wrapper)

    starlette_app = Starlette(
        debug=False,
        routes=[Mount("/mcp", app=handle_mcp)],
        lifespan=lifespan,
    )

    # CORS — only the explicitly allowed origins may make cross-origin
    # requests.  Native MCP clients (non-browser) are unaffected.
    cors_app = CORSMiddleware(
        starlette_app,
        allow_origins=allowed_origins,
        allow_methods=["GET", "POST", "DELETE"],
        allow_headers=[
            "Content-Type",
            "Accept",
            "Mcp-Session-Id",
            "Last-Event-Id",
            "Mcp-Protocol-Version",
        ],
        expose_headers=["Mcp-Session-Id"],
    )

    # Wrap: body-size check → error sanitiser → CORS → Starlette
    secured_app = _BodySizeLimitMiddleware(_ErrorSanitiserMiddleware(cors_app))

    uvicorn.run(
        secured_app,  # type: ignore[arg-type]
        host=host,
        port=port,
        log_level="warning",
    )


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        prog="eip-mcp",
        description="MCP server for the Exploit Intelligence Platform",
    )
    p.add_argument(
        "--version", "-V",
        action="version",
        version=f"eip-mcp {__version__}",
    )
    p.add_argument(
        "--transport",
        choices=["stdio", "streamable-http"],
        default=os.environ.get("EIP_MCP_TRANSPORT", "stdio"),
        help=(
            "Transport protocol (default: stdio).  "
            "Env: EIP_MCP_TRANSPORT"
        ),
    )
    p.add_argument(
        "--host",
        default=os.environ.get("EIP_MCP_HOST", "127.0.0.1"),
        help=(
            "Bind address for HTTP transport (default: 127.0.0.1).  "
            "Env: EIP_MCP_HOST"
        ),
    )
    p.add_argument(
        "--port",
        type=int,
        default=_int_env("EIP_MCP_PORT", 8080),
        help=(
            "Bind port for HTTP transport (default: 8080).  "
            "Env: EIP_MCP_PORT"
        ),
    )
    p.add_argument(
        "--stateless",
        action="store_true",
        default=os.environ.get("EIP_MCP_STATELESS", "").lower()
        in ("1", "true", "yes"),
        help=(
            "Disable session tracking (enables horizontal scaling).  "
            "Env: EIP_MCP_STATELESS"
        ),
    )
    return p.parse_args()


def main():
    """Entry point for the MCP server."""
    args = _parse_args()

    if args.transport == "streamable-http":
        _run_http_server(
            host=args.host,
            port=args.port,
            stateless=args.stateless,
        )
    else:
        asyncio.run(_run_stdio())


if __name__ == "__main__":
    main()
