websafe-kb/scripts/intel/sources/nvd_api.py

from __future__ import annotations

import os
import threading
import time
from hashlib import sha1
from pathlib import Path
from typing import Any, Dict, List

import requests

from intel.http_client import request
from intel.models import Candidate
from intel.config import STATE_DIR
from intel.utils import isoformat, now_utc, parse_dt, read_json, unique, write_json


API_URL = "https://services.nvd.nist.gov/rest/json/cves/2.0"
PUBLIC_INTERVAL_SECONDS = 7.0
DEFAULT_CACHE_TTL_SECONDS = 6 * 60 * 60
_NVD_RATE_LOCK = threading.Lock()
_NVD_LAST_REQUEST = 0.0
_CACHE_DIR = STATE_DIR / "cache" / "nvd"


def _wait_for_slot() -> None:
    global _NVD_LAST_REQUEST
    if os.environ.get("NVD_API_KEY"):
        return
    with _NVD_RATE_LOCK:
        elapsed = time.monotonic() - _NVD_LAST_REQUEST
        if elapsed < PUBLIC_INTERVAL_SECONDS:
            time.sleep(PUBLIC_INTERVAL_SECONDS - elapsed)
        _NVD_LAST_REQUEST = time.monotonic()


def request_nvd(source: Dict[str, Any], headers: Dict[str, Any], params: Dict[str, Any]) -> requests.Response:
    _wait_for_slot()
    response = request("GET", API_URL, source=source, headers=headers, params=params)
    if response.status_code == 429 and not os.environ.get("NVD_API_KEY"):
        time.sleep(PUBLIC_INTERVAL_SECONDS)
        _wait_for_slot()
        response = request("GET", API_URL, source=source, headers=headers, params=params)
    return response


def _cache_ttl_seconds() -> int:
    configured = os.environ.get("WEBSAFE_NVD_CACHE_TTL_SECONDS")
    if configured:
        try:
            return max(0, int(configured))
        except ValueError:
            return DEFAULT_CACHE_TTL_SECONDS
    return DEFAULT_CACHE_TTL_SECONDS


def _cache_key(params: Dict[str, Any]) -> str:
    normalized = "&".join(f"{key}={params[key]}" for key in sorted(params))
    return sha1(normalized.encode("utf-8")).hexdigest()


def _cache_path(params: Dict[str, Any]) -> Path:
    return _CACHE_DIR / f"{_cache_key(params)}.json"


def _load_cached_payload(params: Dict[str, Any]) -> Dict[str, Any] | None:
    ttl_seconds = _cache_ttl_seconds()
    if ttl_seconds <= 0:
        return None
    path = _cache_path(params)
    cached = read_json(path, default=None)
    if not isinstance(cached, dict):
        return None
    fetched_at = parse_dt(cached.get("fetched_at"))
    if fetched_at is None:
        return None
    age = (now_utc() - fetched_at).total_seconds()
    if age > ttl_seconds:
        return None
    payload = cached.get("payload")
    return payload if isinstance(payload, dict) else None


def _write_cached_payload(params: Dict[str, Any], payload: Dict[str, Any]) -> None:
    write_json(
        _cache_path(params),
        {
            "fetched_at": isoformat(now_utc()),
            "payload": payload,
        },
    )


def request_nvd_json(source: Dict[str, Any], headers: Dict[str, Any], params: Dict[str, Any]) -> Dict[str, Any]:
    api_key = os.environ.get("NVD_API_KEY")
    if not api_key:
        cached = _load_cached_payload(params)
        if cached is not None:
            return cached
    response = request_nvd(source, headers, params)
    response.raise_for_status()
    payload = response.json()
    if not isinstance(payload, dict):
        raise ValueError("NVD response payload was not an object")
    if not api_key:
        _write_cached_payload(params, payload)
    return payload


def fetch(system: Dict[str, Any], source: Dict[str, Any]) -> List[Candidate]:
    params = {
        "keywordSearch": source.get("keyword") or system["display_name"],
        "resultsPerPage": source.get("results_per_page", 50),
    }
    headers = {"User-Agent": "websafe-intel"}
    api_key = os.environ.get("NVD_API_KEY")
    if api_key:
        headers["apiKey"] = api_key

    payload = request_nvd_json(source, headers, params)

    candidates: List[Candidate] = []
    for item in payload.get("vulnerabilities", []):
        cve = item.get("cve", {})
        descriptions = cve.get("descriptions", [])
        description = next((d.get("value") for d in descriptions if d.get("lang") == "en"), "")
        metrics = cve.get("metrics", {})
        severity = "unknown"
        cvss_score = None
        for key in ("cvssMetricV31", "cvssMetricV30", "cvssMetricV2"):
            entries = metrics.get(key, [])
            if entries:
                data = entries[0].get("cvssData", {})
                severity = (entries[0].get("baseSeverity") or data.get("baseSeverity") or "unknown").lower()
                cvss_score = data.get("baseScore")
                break

        refs = [ref.get("url") for ref in cve.get("references", []) if ref.get("url")]
        candidates.append(
            Candidate(
                system_id=system["system_id"],
                display_name=system["display_name"],
                category=system["category"],
                advisory_mode=source.get("advisory_mode", "core"),
                source_kind=source["kind"],
                source_name=source["name"],
                source_confidence=source["confidence"],
                source_url=refs[0] if refs else API_URL,
                title=cve.get("id") or f"NVD advisory for {system['display_name']}",
                published_at=cve.get("published"),
                updated_at=cve.get("lastModified"),
                summary=description or "",
                severity=severity,
                cvss_score=cvss_score,
                aliases=unique([cve.get("id")]),
                cve_ids=[cve.get("id")] if cve.get("id") else [],
                references=refs,
                raw=item,
            )
        )
    return candidates