#!/usr/bin/env python3
"""Comprehensive MCP tool integration tests.

Tests all available tools via the Streamable HTTP transport against the live EIP API.
Run: python3 test_mcp_tools.py [--base-url http://host:port/mcp/]
"""

import json
import sys
import time
import requests
import argparse
import re

# ─── Config ───────────────────────────────────────────────────────────────

DEFAULT_BASE_URL = "http://127.0.0.1:8080/mcp/"
HEADERS = {
    "Content-Type": "application/json",
    "Accept": "text/event-stream, application/json",
}

# ─── Helpers ──────────────────────────────────────────────────────────────

_id_counter = 0
_session_id = None
_pass = 0
_fail = 0
_errors = []
_MAX_RATE_LIMIT_RETRIES = 6


def _next_id():
    global _id_counter
    _id_counter += 1
    return _id_counter


def _send(method, params=None, base_url=DEFAULT_BASE_URL):
    """Send a JSON-RPC request and parse the SSE response."""
    global _session_id
    payload = {
        "jsonrpc": "2.0",
        "id": _next_id(),
        "method": method,
        "params": params or {},
    }
    hdrs = dict(HEADERS)
    if _session_id:
        hdrs["Mcp-Session-Id"] = _session_id

    resp = requests.post(base_url, json=payload, headers=hdrs, timeout=120, stream=True)

    # Capture session ID from response
    if "Mcp-Session-Id" in resp.headers:
        _session_id = resp.headers["Mcp-Session-Id"]

    if resp.status_code != 200:
        return {"error": f"HTTP {resp.status_code}: {resp.text[:200]}"}

    # Parse SSE response — handle multi-line data fields.
    # SSE spec: events are separated by blank lines. Data lines start with "data: ".
    # Multiple "data:" lines in the same event are concatenated with newlines.
    result = None
    raw = resp.text
    # Split by double newline to get individual events
    events = raw.split("\n\n")
    for event in events:
        data_parts = []
        for line in event.split("\n"):
            if line.startswith("data: "):
                data_parts.append(line[6:])
            elif line.startswith("data:"):
                data_parts.append(line[5:])
        if data_parts:
            full_data = "\n".join(data_parts)
            try:
                parsed = json.loads(full_data)
                if "result" in parsed:
                    result = parsed["result"]
                elif "error" in parsed:
                    result = {"error": parsed["error"]}
            except json.JSONDecodeError:
                continue
    return result


def _retry_after_seconds(message: str) -> float | None:
    """Extract retry-after seconds from a rate-limit message when present."""
    m = re.search(r"retry after\s+([0-9]+(?:\.[0-9]+)?)s", message, flags=re.IGNORECASE)
    if not m:
        return None
    try:
        return float(m.group(1))
    except ValueError:
        return None


def _is_rate_limited(message: str) -> bool:
    lowered = message.lower()
    return "rate limited locally" in lowered or "rate limited by the eip api" in lowered


def call_tool(name, arguments=None, base_url=DEFAULT_BASE_URL):
    """Call an MCP tool and return (text, error_text).

    Returns (text, None) for success, (None, error_text) for MCP-level errors,
    and (text, text) for tool-level isError=True responses (validation errors etc).
    """
    for attempt in range(1, _MAX_RATE_LIMIT_RETRIES + 1):
        result = _send("tools/call", {"name": name, "arguments": arguments or {}}, base_url)
        if result is None:
            return None, "No response received"
        if isinstance(result, dict) and "error" in result:
            err = str(result["error"])
            if attempt < _MAX_RATE_LIMIT_RETRIES and _is_rate_limited(err):
                wait_s = (_retry_after_seconds(err) or 1.0) + 0.2
                time.sleep(wait_s)
                continue
            return None, err
        # Extract text from content
        content = result.get("content", [])
        is_error = result.get("isError", False)
        text = "\n".join(c.get("text", "") for c in content if c.get("type") == "text")
        if is_error:
            if attempt < _MAX_RATE_LIMIT_RETRIES and _is_rate_limited(text):
                wait_s = (_retry_after_seconds(text) or 1.0) + 0.2
                time.sleep(wait_s)
                continue
            return text, text  # Return text in both slots for error responses
        return text, None
    return None, "Rate-limited after retries"


def check(test_name, text, error, *assertions):
    """Check test result and record pass/fail."""
    global _pass, _fail
    if error:
        print(f"  FAIL  {test_name}: {error[:120]}")
        _fail += 1
        _errors.append((test_name, error[:200]))
        return False
    if text is None:
        print(f"  FAIL  {test_name}: No text response")
        _fail += 1
        _errors.append((test_name, "No text"))
        return False

    for assertion_name, condition in assertions:
        if not condition:
            snippet = text[:150].replace("\n", " ")
            print(f"  FAIL  {test_name} [{assertion_name}]: got: {snippet}")
            _fail += 1
            _errors.append((test_name, f"Assertion failed: {assertion_name}"))
            return False

    print(f"  PASS  {test_name}")
    _pass += 1
    return True


# ─── Test suites ──────────────────────────────────────────────────────────

def test_initialize(base_url):
    """Test MCP initialization."""
    print("\n=== INITIALIZATION ===")
    result = _send("initialize", {
        "protocolVersion": "2025-03-26",
        "capabilities": {},
        "clientInfo": {"name": "test-suite", "version": "1.0"},
    }, base_url)
    ok = result and "serverInfo" in result
    check("initialize", json.dumps(result or {}), None,
          ("has serverInfo", ok),
          ("version", bool(result.get("serverInfo", {}).get("version")) if result else False))

    # Send initialized notification
    _send("notifications/initialized", {}, base_url)


def test_list_tools(base_url):
    """Test tool listing."""
    print("\n=== LIST TOOLS ===")
    result = _send("tools/list", {}, base_url)
    tools = result.get("tools", []) if result else []
    names = [t["name"] for t in tools]
    check("list_tools", json.dumps({"count": len(tools)}), None,
          ("has >=16 tools", len(tools) >= 16),
          ("has search_vulnerabilities", "search_vulnerabilities" in names),
          ("has audit_stack", "audit_stack" in names),
          ("has generate_finding", "generate_finding" in names))
    return names


def test_search_vulnerabilities(base_url):
    """Test vulnerability search with various filters."""
    print("\n=== SEARCH VULNERABILITIES ===")

    # Basic text search
    text, err = call_tool("search_vulnerabilities", {"query": "log4j"}, base_url)
    check("search_vulns: log4j", text, err,
          ("found results", "Found" in (text or "")),
          ("has CVE", "CVE-" in (text or "")))

    # Severity filter
    text, err = call_tool("search_vulnerabilities", {"severity": "critical", "has_exploits": True, "per_page": 5}, base_url)
    check("search_vulns: critical+exploits", text, err,
          ("found results", "Found" in (text or "")),
          ("has CRITICAL", "CRITICAL" in (text or "")))

    # Vendor filter
    text, err = call_tool("search_vulnerabilities", {"vendor": "fortinet", "is_kev": True}, base_url)
    check("search_vulns: fortinet+kev", text, err,
          ("found results", "Found" in (text or "")),
          ("has KEV", "KEV" in (text or "")))

    # CVSS/EPSS thresholds
    text, err = call_tool("search_vulnerabilities", {"min_cvss": 9.0, "min_epss": 0.9, "per_page": 3}, base_url)
    check("search_vulns: high cvss+epss", text, err,
          ("found results", "Found" in (text or "")))

    # Ecosystem filter
    text, err = call_tool("search_vulnerabilities", {"ecosystem": "npm", "has_exploits": True, "per_page": 3}, base_url)
    check("search_vulns: npm ecosystem", text, err,
          ("found results", "Found" in (text or "") or "No results" in (text or "")))

    # CWE filter
    text, err = call_tool("search_vulnerabilities", {"cwe": "79", "severity": "high", "per_page": 3}, base_url)
    check("search_vulns: CWE-79 XSS", text, err,
          ("found results", "Found" in (text or "")))

    # Sort + pagination
    text, err = call_tool("search_vulnerabilities", {"vendor": "microsoft", "sort": "epss_desc", "page": 2, "per_page": 5}, base_url)
    check("search_vulns: sort+page", text, err,
          ("page 2", "page 2/" in (text or "")))

    # Year filter
    text, err = call_tool("search_vulnerabilities", {"year": 2024, "severity": "critical", "per_page": 3}, base_url)
    check("search_vulns: year 2024", text, err,
          ("found results", "Found" in (text or "")))

    # Date range
    text, err = call_tool("search_vulnerabilities", {"date_from": "2025-01-01", "date_to": "2025-06-30", "has_exploits": True, "per_page": 3}, base_url)
    check("search_vulns: date range", text, err,
          ("found results", "Found" in (text or "")))

    # Nuclei filter
    text, err = call_tool("search_vulnerabilities", {"has_nuclei": True, "severity": "critical", "per_page": 3}, base_url)
    check("search_vulns: nuclei templates", text, err,
          ("found results", "Found" in (text or "")))

    # Product filter
    text, err = call_tool("search_vulnerabilities", {"vendor": "apache", "product": "http_server", "per_page": 3}, base_url)
    check("search_vulns: apache/http_server", text, err,
          ("found results", "Found" in (text or "")))


def test_get_vulnerability(base_url):
    """Test vulnerability detail retrieval."""
    print("\n=== GET VULNERABILITY ===")

    # Well-known CVE
    text, err = call_tool("get_vulnerability", {"cve_id": "CVE-2024-3400"}, base_url)
    check("get_vuln: CVE-2024-3400", text, err,
          ("has title", "Title:" in (text or "")),
          ("has CVSS", "CVSS:" in (text or "")),
          ("has DESCRIPTION", "DESCRIPTION:" in (text or "")),
          ("has EXPLOITS", "EXPLOITS" in (text or "")),
          ("is critical", "CRITICAL" in (text or "")))

    # Log4Shell
    text, err = call_tool("get_vulnerability", {"cve_id": "CVE-2021-44228"}, base_url)
    check("get_vuln: log4shell", text, err,
          ("has KEV", "KEV" in (text or "")),
          ("has EPSS", "EPSS:" in (text or "")))

    # A recent CVE
    text, err = call_tool("get_vulnerability", {"cve_id": "CVE-2024-21887"}, base_url)
    check("get_vuln: CVE-2024-21887 (Ivanti)", text, err,
          ("has data", "Title:" in (text or "")))


def test_get_exploit_code(base_url):
    """Test exploit code retrieval."""
    print("\n=== GET EXPLOIT CODE ===")

    # First find an exploit for a known CVE
    text, err = call_tool("search_exploits", {"cve": "CVE-2024-3400", "has_code": True, "per_page": 3}, base_url)
    if text and "[id=" in text:
        # Extract first exploit ID
        import re
        match = re.search(r'\[id=(\d+)\]', text)
        if match:
            eid = int(match.group(1))
            text2, err2 = call_tool("get_exploit_code", {"exploit_id": eid}, base_url)
            check(f"get_code: exploit {eid}", text2, err2,
                  ("has content", len(text2 or "") > 50),
                  ("has file path", "File:" in (text2 or "")))
        else:
            check("get_code: parse id", text, "Could not parse exploit ID")
    else:
        check("get_code: find exploit", text, err or "No exploits with code found")


def test_get_nuclei_templates(base_url):
    """Test Nuclei template retrieval."""
    print("\n=== GET NUCLEI TEMPLATES ===")

    text, err = call_tool("get_nuclei_templates", {"cve_id": "CVE-2024-3400"}, base_url)
    check("nuclei: CVE-2024-3400", text, err,
          ("has template", "Template:" in (text or "")),
          ("has nuclei cmd", "nuclei -t" in (text or "")))

    # Try one that might have dorks
    text, err = call_tool("get_nuclei_templates", {"cve_id": "CVE-2024-27198"}, base_url)
    check("nuclei: CVE-2024-27198 (TeamCity)", text, err,
          ("has template or none", "Template:" in (text or "") or "No Nuclei" in (text or "")))


def test_search_exploits(base_url):
    """Test exploit search with various filters."""
    print("\n=== SEARCH EXPLOITS ===")

    # By source
    text, err = call_tool("search_exploits", {"source": "metasploit", "per_page": 5}, base_url)
    check("exploits: metasploit", text, err,
          ("found results", "Found" in (text or "")),
          ("has metasploit", "metasploit" in (text or "").lower()))

    # By language
    text, err = call_tool("search_exploits", {"language": "python", "per_page": 5}, base_url)
    check("exploits: python", text, err,
          ("found results", "Found" in (text or "")))

    # Trojan detection
    text, err = call_tool("search_exploits", {"llm_classification": "trojan", "per_page": 3}, base_url)
    check("exploits: trojans", text, err,
          ("found or none", "Found" in (text or "") or "No exploits" in (text or "")))

    # AI analysis filters
    text, err = call_tool("search_exploits", {"attack_type": "RCE", "reliability": "reliable", "per_page": 5}, base_url)
    check("exploits: reliable RCE", text, err,
          ("found results", "Found" in (text or "")))

    # By vendor
    text, err = call_tool("search_exploits", {"vendor": "fortinet", "per_page": 5}, base_url)
    check("exploits: fortinet vendor", text, err,
          ("found results", "Found" in (text or "")))

    # By CVE
    text, err = call_tool("search_exploits", {"cve": "CVE-2024-3400", "per_page": 5}, base_url)
    check("exploits: by CVE", text, err,
          ("found results", "Found" in (text or "")))

    # By author
    text, err = call_tool("search_exploits", {"author": "Chocapikk", "per_page": 5}, base_url)
    check("exploits: by author", text, err,
          ("response", (text or "") != ""))

    # Complexity filter
    text, err = call_tool("search_exploits", {"complexity": "trivial", "attack_type": "RCE", "per_page": 3}, base_url)
    check("exploits: trivial RCE", text, err,
          ("found or none", (text or "") != ""))

    # Stars filter
    text, err = call_tool("search_exploits", {"min_stars": 100, "sort": "stars_desc", "per_page": 5}, base_url)
    check("exploits: popular (100+ stars)", text, err,
          ("found results", "Found" in (text or "")))


def test_list_authors(base_url):
    """Test author listing."""
    print("\n=== LIST AUTHORS ===")

    text, err = call_tool("list_authors", {"per_page": 10}, base_url)
    check("authors: list", text, err,
          ("has authors", "Exploit Authors" in (text or "")),
          ("has exploit count", "exploits" in (text or "")))

    # Page 2
    text, err = call_tool("list_authors", {"page": 2, "per_page": 10}, base_url)
    check("authors: page 2", text, err,
          ("has authors", "Exploit Authors" in (text or "")))


def test_get_author(base_url):
    """Test author detail."""
    print("\n=== GET AUTHOR ===")

    text, err = call_tool("get_author", {"author_name": "Chocapikk"}, base_url)
    check("author: Chocapikk", text, err,
          ("has author", "Author:" in (text or "")),
          ("has exploits", "Exploits:" in (text or "")))


def test_list_cwes(base_url):
    """Test CWE listing."""
    print("\n=== LIST CWES ===")

    text, err = call_tool("list_cwes", {}, base_url)
    check("cwes: list", text, err,
          ("has CWE data", "CWE Categories" in (text or "")),
          ("has CWE-79", "CWE-79" in (text or "")))


def test_get_cwe(base_url):
    """Test CWE detail."""
    print("\n=== GET CWE ===")

    text, err = call_tool("get_cwe", {"cwe_id": "CWE-79"}, base_url)
    check("cwe: CWE-79", text, err,
          ("has name", "CWE-79:" in (text or "")),
          ("has vuln count", "Vulnerabilities:" in (text or "")))

    # Numeric form
    text, err = call_tool("get_cwe", {"cwe_id": "89"}, base_url)
    check("cwe: 89 (numeric)", text, err,
          ("has SQL injection", "CWE-89:" in (text or "")))


def test_list_vendors(base_url):
    """Test vendor listing."""
    print("\n=== LIST VENDORS ===")

    text, err = call_tool("list_vendors", {}, base_url)
    check("vendors: list", text, err,
          ("has vendors", "Top Vendors" in (text or "")),
          ("has microsoft", "microsoft" in (text or "").lower()))


def test_list_products(base_url):
    """Test product listing for a vendor."""
    print("\n=== LIST PRODUCTS ===")

    text, err = call_tool("list_products", {"vendor": "microsoft"}, base_url)
    check("products: microsoft", text, err,
          ("has products", "Products for" in (text or "")),
          ("has windows", "windows" in (text or "").lower()))

    text, err = call_tool("list_products", {"vendor": "fortinet"}, base_url)
    check("products: fortinet", text, err,
          ("has products", "Products for" in (text or "")))


def test_lookup_alt_id(base_url):
    """Test alternate ID lookup."""
    print("\n=== LOOKUP ALT ID ===")

    text, err = call_tool("lookup_alt_id", {"alt_id": "EDB-48537"}, base_url)
    check("lookup: EDB-48537", text, err,
          ("has CVE or not found", "CVE-" in (text or "") or "No CVE" in (text or "")))

    text, err = call_tool("lookup_alt_id", {"alt_id": "GHSA-jfh8-c2jp-5v3q"}, base_url)
    check("lookup: GHSA", text, err,
          ("response", (text or "") != ""))


def test_audit_stack(base_url):
    """Test technology stack auditing."""
    print("\n=== AUDIT STACK ===")

    text, err = call_tool("audit_stack", {"technologies": "nginx, postgresql, node.js"}, base_url)
    check("audit: nginx,pg,node", text, err,
          ("has header", "STACK AUDIT" in (text or "")),
          ("has nginx", "NGINX" in (text or "")),
          ("has postgresql", "POSTGRESQL" in (text or "")))

    # Single tech
    text, err = call_tool("audit_stack", {"technologies": "apache"}, base_url)
    check("audit: apache", text, err,
          ("has results", "STACK AUDIT" in (text or "")))


def test_generate_finding(base_url):
    """Test pentest finding generation."""
    print("\n=== GENERATE FINDING ===")

    text, err = call_tool("generate_finding", {
        "cve_id": "CVE-2024-3400",
        "target": "fw.corp.example.com",
        "notes": "Confirmed RCE via GlobalProtect gateway during external assessment",
    }, base_url)
    check("finding: CVE-2024-3400", text, err,
          ("has markdown header", "# CVE-2024-3400" in (text or "")),
          ("has severity", "Severity:" in (text or "")),
          ("has CVSS", "CVSS" in (text or "")),
          ("has target", "fw.corp.example.com" in (text or "")),
          ("has notes", "Tester Notes" in (text or "")),
          ("has exploit section", "Exploit Availability" in (text or "")))

    # Without target/notes
    text, err = call_tool("generate_finding", {"cve_id": "CVE-2021-44228"}, base_url)
    check("finding: log4shell (minimal)", text, err,
          ("has markdown", "# CVE-2021-44228" in (text or "")))


def test_platform_stats(base_url):
    """Test platform stats."""
    print("\n=== PLATFORM STATS ===")

    text, err = call_tool("get_platform_stats", {}, base_url)
    check("stats", text, err,
          ("has title", "Exploit Intelligence Platform" in (text or "")),
          ("has vuln count", "Total Vulnerabilities:" in (text or "")),
          ("has exploit count", "Total Exploits:" in (text or "")))


def test_check_health(base_url):
    """Test health check."""
    print("\n=== CHECK HEALTH ===")

    text, err = call_tool("check_health", {}, base_url)
    check("health", text, err,
          ("has status", "API Status:" in (text or "")),
          ("healthy", "healthy" in (text or "").lower() or "ok" in (text or "").lower()),
          ("has sources", "Ingestion Sources:" in (text or "")))


# ─── Validation / Error tests ────────────────────────────────────────────

def test_validation_errors(base_url):
    """Test input validation rejects bad input."""
    print("\n=== VALIDATION & ERROR HANDLING ===")

    def check_err(name, text, err, *keywords):
        """Validation tests expect an error. Check that error contains keywords."""
        combined = (err or text or "")
        found = any(kw.lower() in combined.lower() for kw in keywords)
        global _pass, _fail
        if found:
            print(f"  PASS  {name}")
            _pass += 1
            return True
        else:
            snippet = combined[:120].replace("\n", " ")
            print(f"  FAIL  {name}: unexpected: {snippet}")
            _fail += 1
            _errors.append((name, snippet))
            return False

    # Empty search (no query, no filters)
    text, err = call_tool("search_vulnerabilities", {}, base_url)
    check_err("validation: empty search", text, err, "Provide a search", "Validation")

    # Invalid CVE format
    text, err = call_tool("get_vulnerability", {"cve_id": "not-a-cve"}, base_url)
    check_err("validation: bad CVE format", text, err, "Invalid", "Validation")

    # Invalid severity
    text, err = call_tool("search_vulnerabilities", {"severity": "superduper", "query": "test"}, base_url)
    check_err("validation: bad severity", text, err, "Invalid", "Validation", "not one of")

    # CVSS out of range
    text, err = call_tool("search_vulnerabilities", {"min_cvss": 15.0, "query": "test"}, base_url)
    check_err("validation: cvss > 10", text, err, "must be", "Validation")

    # EPSS out of range
    text, err = call_tool("search_vulnerabilities", {"min_epss": 2.0, "query": "test"}, base_url)
    check_err("validation: epss > 1", text, err, "must be", "Validation")

    # per_page too large
    text, err = call_tool("search_vulnerabilities", {"query": "test", "per_page": 100}, base_url)
    check_err("validation: per_page > 25", text, err, "must be", "Validation")

    # Path traversal in file_path
    text, err = call_tool("get_exploit_code", {"exploit_id": 1, "file_path": "../secret.txt"}, base_url)
    check_err("validation: path traversal", text, err, "traversal", "Validation")

    # Null byte injection
    text, err = call_tool("search_vulnerabilities", {"query": "test\x00inject"}, base_url)
    check_err("validation: null byte", text, err, "Null byte", "Validation")

    # 404 - nonexistent CVE
    text, err = call_tool("get_vulnerability", {"cve_id": "CVE-9999-99999"}, base_url)
    check_err("error: 404 nonexistent CVE", text, err, "Not found", "not found")

    # Unknown tool
    text, err = call_tool("nonexistent_tool", {}, base_url)
    check_err("error: unknown tool", text, err, "Unknown", "error", "not found")

    # Invalid ecosystem
    text, err = call_tool("search_vulnerabilities", {"ecosystem": "fake_ecosystem", "query": "test"}, base_url)
    check_err("validation: bad ecosystem", text, err, "Invalid", "Validation", "not one of")

    # Too many technologies
    text, err = call_tool("audit_stack", {"technologies": "a, b, c, d, e, f"}, base_url)
    check_err("validation: >5 technologies", text, err, "Maximum", "Validation")

    # Invalid exploit source
    text, err = call_tool("search_exploits", {"source": "fake_source"}, base_url)
    check_err("validation: bad source", text, err, "Invalid", "Validation")

    # Negative exploit ID
    text, err = call_tool("get_exploit_code", {"exploit_id": -5}, base_url)
    check_err("validation: negative exploit_id", text, err, "Validation", "range")

    # Author not found
    text, err = call_tool("get_author", {"author_name": "zzz_nonexistent_author_xyz"}, base_url)
    check_err("error: author 404", text, err, "not found")

    # CWE not found
    text, err = call_tool("get_cwe", {"cwe_id": "99999"}, base_url)
    check_err("error: CWE 404", text, err, "not found")

    # Vendor with no products
    text, err = call_tool("list_products", {"vendor": "zzz_nonexistent_vendor_xyz"}, base_url)
    check_err("error: vendor no products", text, err, "No products", "not found")


def test_edge_cases(base_url):
    """Test boundary/edge case scenarios."""
    print("\n=== EDGE CASES ===")

    # Unicode in query
    text, err = call_tool("search_vulnerabilities", {"query": "apache \u2014 httpd"}, base_url)
    check("edge: unicode dash in query", text, err,
          ("handled", text is not None))

    # Unicode dashes in CVE ID
    text, err = call_tool("get_vulnerability", {"cve_id": "CVE\u20132024\u20133400"}, base_url)
    check("edge: unicode dashes in CVE", text, err,
          ("normalized", "CVE-2024-3400" in (text or "") or "Title:" in (text or "")))

    # Case insensitive CVE
    text, err = call_tool("get_vulnerability", {"cve_id": "cve-2024-3400"}, base_url)
    check("edge: lowercase CVE", text, err,
          ("normalized", "Title:" in (text or "")))

    # CWE with prefix
    text, err = call_tool("get_cwe", {"cwe_id": "CWE-79"}, base_url)
    check("edge: CWE-79 with prefix", text, err,
          ("works", "CWE-79:" in (text or "")))

    # Max per_page (25)
    text, err = call_tool("search_vulnerabilities", {"query": "apache", "per_page": 25}, base_url)
    check("edge: max per_page=25", text, err,
          ("works", "Found" in (text or "")))

    # Page 1 explicitly
    text, err = call_tool("search_vulnerabilities", {"query": "apache", "page": 1}, base_url)
    check("edge: explicit page 1", text, err,
          ("works", "page 1/" in (text or "")))

    # Audit with single tech that has spaces
    text, err = call_tool("audit_stack", {"technologies": "node.js"}, base_url)
    check("edge: tech with dots", text, err,
          ("works", "STACK AUDIT" in (text or "")))


# ─── Main ─────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
    args = parser.parse_args()
    base_url = args.base_url

    print(f"Testing eip-mcp at {base_url}")
    print("=" * 60)

    start = time.time()

    test_initialize(base_url)
    test_list_tools(base_url)
    test_search_vulnerabilities(base_url)
    test_get_vulnerability(base_url)
    test_get_exploit_code(base_url)
    test_get_nuclei_templates(base_url)
    test_search_exploits(base_url)
    test_list_authors(base_url)
    test_get_author(base_url)
    test_list_cwes(base_url)
    test_get_cwe(base_url)
    test_list_vendors(base_url)
    test_list_products(base_url)
    test_lookup_alt_id(base_url)
    test_audit_stack(base_url)
    test_generate_finding(base_url)
    test_platform_stats(base_url)
    test_check_health(base_url)
    # Pause before validation/edge tests to avoid rate limiting
    print("\n  (pausing 5s to avoid rate limiting...)")
    time.sleep(5)
    test_validation_errors(base_url)
    # Pause again
    print("\n  (pausing 5s to avoid rate limiting...)")
    time.sleep(5)
    test_edge_cases(base_url)

    elapsed = time.time() - start

    print("\n" + "=" * 60)
    print(f"RESULTS: {_pass} passed, {_fail} failed  ({elapsed:.1f}s)")
    print("=" * 60)

    if _errors:
        print("\nFAILURES:")
        for name, msg in _errors:
            print(f"  - {name}: {msg}")

    sys.exit(0 if _fail == 0 else 1)


if __name__ == "__main__":
    main()
