CVE Risk MCP Server

from __future__ import annotations

import json
from typing import Any

from cve_risk_mcp.tools import CveRiskService, handle_tool_call, tool_definitions


class StubCveClient:
    def __init__(self, payload: dict[str, Any]) -> None:
        self.payload = payload

    def fetch(self, cve_id: str) -> dict[str, Any]:
        return dict(self.payload, cveMetadata={"cveId": cve_id})


class StubKevClient:
    def __init__(self, hit: bool = True) -> None:
        self.hit = hit

    def lookup(self, cve_id: str) -> dict[str, Any] | None:
        if self.hit:
            return {"cveID": cve_id, "vendorProject": "Test", "product": "Widget"}
        return None


class StubEpssClient:
    def __init__(self, epss: float = 0.0, percentile: float | None = None) -> None:
        self.epss = epss
        self.percentile = percentile

    def lookup(self, cve_id: str) -> dict[str, Any] | None:
        return {"epss": self.epss, "percentile": self.percentile, "date": "2026-01-31"}


def test_tool_definitions_include_expected_tools() -> None:
    tools = tool_definitions()
    names = {tool["name"] for tool in tools}
    assert {"cve_lookup", "cve_rank", "cve_patch_first", "cve_status", "cve_refresh"}.issubset(
        names
    )


def test_tool_call_lookup() -> None:
    payload: dict[str, Any] = {
        "containers": {
            "cna": {
                "descriptions": [{"lang": "en", "value": "Sample"}],
                "metrics": [
                    {
                        "cvssV3_1": {
                            "version": "3.1",
                            "baseSeverity": "HIGH",
                            "baseScore": 8.8,
                            "vectorString": "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H",
                        }
                    }
                ],
                "references": [{"url": "https://example.com"}],
            }
        }
    }
    service = CveRiskService(StubCveClient(payload), StubKevClient())
    content = handle_tool_call("cve_lookup", {"cve_id": "CVE-2024-0002"}, service)
    assert content and content[0]["type"] == "text"
    parsed = json.loads(content[0]["text"])
    assert parsed["cve_id"] == "CVE-2024-0002"
    assert parsed["risk"]["priority"] in {"high", "critical"}


def test_tool_call_lookup_includes_epss_when_available() -> None:
    payload: dict[str, Any] = {
        "containers": {
            "cna": {
                "descriptions": [{"lang": "en", "value": "Sample"}],
                "metrics": [{"other": {"content": {"other": "low"}}}],
            }
        }
    }
    service = CveRiskService(
        StubCveClient(payload), StubKevClient(hit=False), StubEpssClient(epss=0.91)
    )
    content = handle_tool_call(
        "cve_lookup",
        {"cve_id": "CVE-2024-0003", "include_epss": True, "include_kev": False},
        service,
    )
    parsed = json.loads(content[0]["text"])
    assert parsed["epss"]["epss"] == 0.91
    assert parsed["risk"]["priority"] in {"medium", "high", "critical"}


def test_tool_call_patch_first() -> None:
    payload: dict[str, Any] = {
        "containers": {"cna": {"metrics": [{"other": {"content": {"other": "medium"}}}]}}
    }
    service = CveRiskService(StubCveClient(payload), StubKevClient(hit=True))
    content = handle_tool_call(
        "cve_patch_first",
        {"cve_ids": ["CVE-2024-0004", "CVE-2024-0005"], "top_n": 1},
        service,
    )
    parsed = json.loads(content[0]["text"])
    assert len(parsed["recommendations"]) == 1
    assert len(parsed["ranked"]) == 2


def test_tool_call_status() -> None:
    payload: dict[str, Any] = {
        "containers": {"cna": {"metrics": [{"other": {"content": {"other": "medium"}}}]}}
    }

    class StatusStub:
        def status(self) -> dict[str, Any]:
            return {"sources": {"cve": {"enabled": True, "last_success_iso": None}}}

    service = CveRiskService(StubCveClient(payload), StubKevClient(hit=False))
    service.status_provider = StatusStub()
    content = handle_tool_call("cve_status", {}, service)
    parsed = json.loads(content[0]["text"])
    assert "sources" in parsed
    assert "cve" in parsed["sources"]


def test_tool_call_refresh() -> None:
    payload: dict[str, Any] = {
        "containers": {"cna": {"metrics": [{"other": {"content": {"other": "medium"}}}]}}
    }

    class RefreshKevClient(StubKevClient):
        def refresh(self) -> int:
            return 42

    service = CveRiskService(StubCveClient(payload), RefreshKevClient(hit=False))
    content = handle_tool_call(
        "cve_refresh",
        {"cve_ids": ["CVE-2024-0007"], "include_epss": False, "refresh_kev": True},
        service,
    )
    parsed = json.loads(content[0]["text"])
    assert parsed["cve_warmed"] == 1
    assert parsed["kev_entries"] == 42


def test_tool_call_export_with_markdown() -> None:
    payload: dict[str, Any] = {
        "containers": {
            "cna": {
                "title": "Sample",
                "metrics": [
                    {
                        "cvssV3_1": {
                            "version": "3.1",
                            "baseSeverity": "HIGH",
                            "baseScore": 8.0,
                            "vectorString": "CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:H/I:H/A:H",
                        }
                    }
                ],
            }
        }
    }
    service = CveRiskService(StubCveClient(payload), StubKevClient(hit=False))
    content = handle_tool_call(
        "cve_export",
        {"cve_ids": ["CVE-2024-0101"], "include_markdown": True},
        service,
    )
    parsed = json.loads(content[0]["text"])
    assert "ranked" in parsed
    assert "markdown" in parsed
    assert "CVE-2024-0101" in parsed["markdown"]