CVE Risk MCP Server

from __future__ import annotations

import json
import os
from typing import Any

import requests
from cve_risk_mcp.cli import main


class DummyResponse:
    def __init__(self, payload: Any) -> None:
        self._payload = payload
        self.headers: dict[str, str] = {}
        self.status_code = 200

    def raise_for_status(self) -> None:
        return None

    def json(self) -> Any:
        return self._payload


def test_cli_lookup_emits_json(capsys: Any, monkeypatch: Any) -> None:
    def fake_get(url: str, *_: Any, **__: Any) -> DummyResponse:
        cve_id = url.rstrip("/").split("/")[-1]
        return DummyResponse(
            {
                "cveMetadata": {"cveId": cve_id, "datePublished": "2024-01-01T00:00:00Z"},
                "containers": {
                    "cna": {
                        "descriptions": [{"lang": "en", "value": "Sample"}],
                        "metrics": [{"other": {"content": {"other": "low"}}}],
                    }
                },
            }
        )

    monkeypatch.setattr(requests, "get", fake_get)
    code = main(
        [
            "--cve-url",
            "https://example.invalid/api/cve",
            "--no-epss",
            "lookup",
            "CVE-2024-0001",
            "--no-include-kev",
            "--no-include-epss",
        ]
    )
    assert code == 0
    stdout = capsys.readouterr().out.strip()
    parsed = json.loads(stdout)
    assert parsed["cve_id"] == "CVE-2024-0001"
    assert parsed["risk"]["priority"] in {"low", "medium", "high", "critical"}


def test_cli_patch_first_emits_json(capsys: Any, monkeypatch: Any) -> None:
    def fake_get(url: str, *_: Any, **__: Any) -> DummyResponse:
        cve_id = url.rstrip("/").split("/")[-1]
        return DummyResponse(
            {
                "cveMetadata": {"cveId": cve_id, "datePublished": "2024-01-01T00:00:00Z"},
                "containers": {
                    "cna": {
                        "metrics": [{"other": {"content": {"other": "medium"}}}],
                    }
                },
            }
        )

    monkeypatch.setattr(requests, "get", fake_get)
    code = main(
        [
            "--cve-url",
            "https://example.invalid/api/cve",
            "--no-epss",
            "patch-first",
            "CVE-2024-0002",
            "CVE-2024-0003",
            "--top-n",
            "1",
            "--no-include-kev",
            "--no-include-epss",
        ]
    )
    assert code == 0
    stdout = capsys.readouterr().out.strip()
    parsed = json.loads(stdout)
    assert len(parsed["ranked"]) == 2
    assert len(parsed["recommendations"]) == 1


def test_cli_status_emits_json(capsys: Any) -> None:
    code = main(["--no-epss", "status"])
    assert code == 0
    stdout = capsys.readouterr().out.strip()
    parsed = json.loads(stdout)
    assert "sources" in parsed
    assert "cve" in parsed["sources"]


def test_cli_refresh_emits_json(capsys: Any) -> None:
    code = main(["--no-epss", "refresh", "--no-refresh-kev"])
    assert code == 0
    stdout = capsys.readouterr().out.strip()
    parsed = json.loads(stdout)
    assert "kev_refreshed" in parsed


def test_cli_export_emits_markdown(capsys: Any, monkeypatch: Any) -> None:
    def fake_get(url: str, *_: Any, **__: Any) -> DummyResponse:
        cve_id = url.rstrip("/").split("/")[-1]
        return DummyResponse(
            {
                "cveMetadata": {"cveId": cve_id},
                "containers": {"cna": {"metrics": [{"other": {"content": {"other": "low"}}}]}},
            }
        )

    monkeypatch.setattr(requests, "get", fake_get)
    code = main(["--no-epss", "export", "CVE-2024-0102", "--include-markdown"])
    assert code == 0
    stdout = capsys.readouterr().out.strip()
    parsed = json.loads(stdout)
    assert "markdown" in parsed


def test_cli_output_writes_file(tmp_path: Any, monkeypatch: Any) -> None:
    def fake_get(url: str, *_: Any, **__: Any) -> DummyResponse:
        cve_id = url.rstrip("/").split("/")[-1]
        return DummyResponse(
            {
                "cveMetadata": {"cveId": cve_id},
                "containers": {"cna": {"metrics": [{"other": {"content": {"other": "low"}}}]}},
            }
        )

    monkeypatch.setattr(requests, "get", fake_get)
    output_path = os.fspath(tmp_path / "out.json")
    code = main(
        [
            "--no-epss",
            "--output",
            output_path,
            "lookup",
            "CVE-2024-0202",
            "--no-include-kev",
            "--no-include-epss",
        ]
    )
    assert code == 0
    written = json.loads((tmp_path / "out.json").read_text())
    assert written["cve_id"] == "CVE-2024-0202"


def test_cli_cve_file_combines_ids(tmp_path: Any, monkeypatch: Any, capsys: Any) -> None:
    file_content = "CVE-2024-0401\nCVE-2024-0402\n"
    file_path = tmp_path / "cves.txt"
    file_path.write_text(file_content)

    def fake_get(url: str, *_: Any, **__: Any) -> DummyResponse:
        cve_id = url.rstrip("/").split("/")[-1]
        return DummyResponse(
            {
                "cveMetadata": {"cveId": cve_id},
                "containers": {"cna": {"metrics": [{"other": {"content": {"other": "low"}}}]}},
            }
        )

    monkeypatch.setattr(requests, "get", fake_get)
    code = main(
        [
            "--no-epss",
            "--format",
            "ndjson",
            "--cve-file",
            str(file_path),
            "rank",
            "CVE-2024-0403",
            "--no-include-kev",
            "--no-include-epss",
        ]
    )
    assert code == 0
    stdout = capsys.readouterr().out.strip()
    lines = stdout.splitlines()
    assert len(lines) == 3


def test_cli_ndjson_format_outputs_lines(capsys: Any, monkeypatch: Any) -> None:
    def fake_get(url: str, *_: Any, **__: Any) -> DummyResponse:
        cve_id = url.rstrip("/").split("/")[-1]
        return DummyResponse(
            {
                "cveMetadata": {"cveId": cve_id},
                "containers": {"cna": {"metrics": [{"other": {"content": {"other": "low"}}}]}},
            }
        )

    monkeypatch.setattr(requests, "get", fake_get)
    code = main(
        [
            "--no-epss",
            "--format",
            "ndjson",
            "rank",
            "CVE-2024-0301",
            "CVE-2024-0302",
            "--no-include-kev",
            "--no-include-epss",
        ]
    )
    assert code == 0
    stdout = capsys.readouterr().out.strip()
    lines = stdout.splitlines()
    assert len(lines) == 2
    parsed = json.loads(lines[0])
    assert "cve_id" in parsed


def test_cli_emits_json_error_and_nonzero_on_bad_payload(capsys: Any, monkeypatch: Any) -> None:
    def fake_get(_: str, *__: Any, **___: Any) -> DummyResponse:
        return DummyResponse([])

    monkeypatch.setattr(requests, "get", fake_get)
    code = main(
        [
            "--cve-url",
            "https://example.invalid/api/cve",
            "--no-epss",
            "lookup",
            "CVE-2024-0009",
            "--no-include-kev",
            "--no-include-epss",
        ]
    )
    assert code == 1
    stdout = capsys.readouterr().out.strip()
    parsed = json.loads(stdout)
    assert parsed["error"]["type"] in {"UpstreamResponseError", "ValueError", "RuntimeError"}