vulncheck/app/services/enrichment_service.py

"""
Vulnerability Enrichment Service

Reichert CVE-Daten mit Threat-Intelligence aus offenen Quellen an:
- EPSS (Exploit Prediction Scoring System) via FIRST.org
- CISA KEV (Known Exploited Vulnerabilities) catalog
- ENISA EUVD (EU exploited + critical vulnerabilities)

Alle Quellen sind kostenlos und benötigen keinen API-Key.
"""
import json
import logging
import re
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Iterable

import httpx
from sqlalchemy.orm import Session

from app.models.setting import Setting
from app.models.vulnerability import Vulnerability

logger = logging.getLogger(__name__)


# ---------- Endpoints ----------
EPSS_API_URL = "https://api.first.org/data/v1/epss"
KEV_FEED_URL = "https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json"
# Full EUVD-exploited catalog (paginated, ~1500+ CVEs).
# The bare /exploitedvulnerabilities endpoint only returns a curated "recent 4" widget.
EUVD_SEARCH_EXPLOITED_URL = "https://euvdservices.enisa.europa.eu/api/search?exploited=true"
# Small curated ENISA-flagged critical list (handful of entries).
EUVD_CRITICAL_URL = "https://euvdservices.enisa.europa.eu/api/criticalvulnerabilities"
# Pagination size for search API
EUVD_PAGE_SIZE = 100
# Hard cap on pages to avoid runaway loop
EUVD_MAX_PAGES = 50

# ---------- Cache keys (settings table) ----------
KEV_CACHE_KEY = "enrichment_kev_cache"
KEV_CACHE_TS_KEY = "enrichment_kev_cache_updated_at"
KEV_TTL_HOURS = 24
EUVD_CACHE_KEY = "enrichment_euvd_cache"
EUVD_CACHE_TS_KEY = "enrichment_euvd_cache_updated_at"
EUVD_TTL_HOURS = 24

# ---------- Toggles ----------
SETTING_EPSS_ENABLED = "enrichment_epss_enabled"
SETTING_KEV_ENABLED = "enrichment_kev_enabled"
SETTING_EUVD_ENABLED = "enrichment_euvd_enabled"

# ---------- HTTP ----------
HTTP_TIMEOUT = 30.0
EPSS_BATCH_SIZE = 100  # FIRST EPSS API supports comma-separated CVE list

# ---------- Helpers ----------
CVE_REGEX = re.compile(r"CVE-\d{4}-\d{4,}")


class EnrichmentError(Exception):
    """Raised when enrichment cannot complete."""


def _setting_bool(db: Session, key: str, default: bool = True) -> bool:
    s = db.query(Setting).filter(Setting.key == key).first()
    if not s or s.value is None:
        return default
    return str(s.value).strip().lower() in {"1", "true", "yes", "on"}


def _set_setting(db: Session, key: str, value: str, description: str = "") -> None:
    s = db.query(Setting).filter(Setting.key == key).first()
    if s:
        s.value = value
    else:
        s = Setting(key=key, value=value, description=description)
        db.add(s)
    db.commit()


# ============================================================
# EPSS
# ============================================================

def fetch_epss_scores(cve_ids: Iterable[str]) -> Dict[str, Dict[str, float]]:
    """
    Fetch EPSS scores for a list of CVE IDs.

    Returns: { "CVE-XXXX-NNNN": {"epss": 0.123, "percentile": 0.94}, ... }
    Missing CVEs are simply absent from the response.
    """
    cve_list = [c for c in cve_ids if c and c.startswith("CVE-")]
    if not cve_list:
        return {}

    results: Dict[str, Dict[str, float]] = {}

    with httpx.Client(timeout=HTTP_TIMEOUT) as client:
        for i in range(0, len(cve_list), EPSS_BATCH_SIZE):
            batch = cve_list[i:i + EPSS_BATCH_SIZE]
            params = {"cve": ",".join(batch)}
            try:
                resp = client.get(EPSS_API_URL, params=params)
                resp.raise_for_status()
                payload = resp.json()
            except httpx.HTTPError as e:
                logger.warning(f"EPSS batch fetch failed ({len(batch)} CVEs): {e}")
                continue

            for entry in payload.get("data", []):
                cve = entry.get("cve")
                if not cve:
                    continue
                try:
                    score = float(entry.get("epss")) if entry.get("epss") is not None else None
                    pct = float(entry.get("percentile")) if entry.get("percentile") is not None else None
                except (TypeError, ValueError):
                    continue
                if score is None:
                    continue
                results[cve] = {"epss": score, "percentile": pct or 0.0}

    logger.info(f"EPSS: fetched {len(results)} scores for {len(cve_list)} CVEs")
    return results


# ============================================================
# CISA KEV
# ============================================================

def _load_kev_cache(db: Session) -> Optional[Dict[str, dict]]:
    """Returns {cve_id: kev_entry} from settings cache, or None if missing/stale."""
    ts_setting = db.query(Setting).filter(Setting.key == KEV_CACHE_TS_KEY).first()
    cache_setting = db.query(Setting).filter(Setting.key == KEV_CACHE_KEY).first()

    if not ts_setting or not cache_setting or not cache_setting.value:
        return None

    try:
        ts = datetime.fromisoformat(ts_setting.value)
    except ValueError:
        return None

    if datetime.now() - ts > timedelta(hours=KEV_TTL_HOURS):
        return None

    try:
        return json.loads(cache_setting.value)
    except json.JSONDecodeError:
        return None


def _store_kev_cache(db: Session, kev_map: Dict[str, dict]) -> None:
    _set_setting(db, KEV_CACHE_KEY, json.dumps(kev_map), "CISA KEV cache (24h TTL)")
    _set_setting(db, KEV_CACHE_TS_KEY, datetime.now().isoformat(), "Timestamp of last KEV refresh")


def fetch_kev_catalog(db: Session, force_refresh: bool = False) -> Dict[str, dict]:
    """
    Fetch CISA KEV catalog. Cached for 24h in `settings` table.

    Returns: { "CVE-XXXX-NNNN": {date_added, ransomware_use, short_description}, ... }
    """
    if not force_refresh:
        cached = _load_kev_cache(db)
        if cached is not None:
            logger.debug(f"KEV: using cached catalog ({len(cached)} entries)")
            return cached

    logger.info("KEV: fetching fresh catalog from CISA")
    try:
        with httpx.Client(timeout=HTTP_TIMEOUT) as client:
            resp = client.get(KEV_FEED_URL)
            resp.raise_for_status()
            payload = resp.json()
    except httpx.HTTPError as e:
        logger.error(f"KEV fetch failed: {e}")
        # Fall back to stale cache if present
        cache_setting = db.query(Setting).filter(Setting.key == KEV_CACHE_KEY).first()
        if cache_setting and cache_setting.value:
            try:
                logger.warning("KEV: returning stale cache after fetch failure")
                return json.loads(cache_setting.value)
            except json.JSONDecodeError:
                pass
        raise EnrichmentError(f"KEV fetch failed and no cache available: {e}")

    kev_map: Dict[str, dict] = {}
    for entry in payload.get("vulnerabilities", []):
        cve = entry.get("cveID")
        if not cve:
            continue
        kev_map[cve] = {
            "date_added": entry.get("dateAdded"),
            "ransomware_use": (entry.get("knownRansomwareCampaignUse") or "").lower() == "known",
            "short_description": entry.get("shortDescription"),
        }

    _store_kev_cache(db, kev_map)
    logger.info(f"KEV: cached {len(kev_map)} entries from CISA")
    return kev_map


# ============================================================
# ENISA EUVD
# ============================================================

def _load_euvd_cache(db: Session) -> Optional[Dict[str, dict]]:
    """Returns {cve_id: euvd_entry} from settings cache, or None if missing/stale."""
    ts_setting = db.query(Setting).filter(Setting.key == EUVD_CACHE_TS_KEY).first()
    cache_setting = db.query(Setting).filter(Setting.key == EUVD_CACHE_KEY).first()

    if not ts_setting or not cache_setting or not cache_setting.value:
        return None

    try:
        ts = datetime.fromisoformat(ts_setting.value)
    except ValueError:
        return None

    if datetime.now() - ts > timedelta(hours=EUVD_TTL_HOURS):
        return None

    try:
        return json.loads(cache_setting.value)
    except json.JSONDecodeError:
        return None


def _store_euvd_cache(db: Session, euvd_map: Dict[str, dict]) -> None:
    _set_setting(db, EUVD_CACHE_KEY, json.dumps(euvd_map), "ENISA EUVD cache (24h TTL)")
    _set_setting(db, EUVD_CACHE_TS_KEY, datetime.now().isoformat(), "Timestamp of last EUVD refresh")


def _parse_euvd_date(raw) -> Optional[str]:
    """
    Convert ENISA date strings to ISO-8601 if possible.
    Returns the original string when we can't parse (caller stores raw).
    ENISA observed formats: 'Apr 29, 2026, 3:10:37 PM', ISO-8601, plain dates.
    """
    if not raw or not isinstance(raw, str):
        return None
    raw = raw.strip()
    # Try ISO first (with possible trailing Z)
    try:
        return datetime.fromisoformat(raw.replace("Z", "+00:00")).isoformat()
    except ValueError:
        pass
    # ENISA's "%b %d, %Y, %I:%M:%S %p" e.g. "Apr 29, 2026, 3:10:37 PM"
    for fmt in ("%b %d, %Y, %I:%M:%S %p", "%b %d, %Y", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d"):
        try:
            return datetime.strptime(raw, fmt).isoformat()
        except ValueError:
            continue
    return raw  # store as-is, downstream handles failure gracefully


def _parse_euvd_entry(entry: dict, is_critical: bool) -> Iterable[tuple]:
    """
    EUVD entries carry one EUVD-ID and CVE aliases (sometimes as
    newline-separated string, sometimes as list).
    Yields (cve_id, payload) tuples for every CVE we can extract.
    """
    euvd_id = entry.get("id") or entry.get("euvd_id") or entry.get("euvdId")

    # Prefer exploitedSince > datePublished > others
    raw_date = (
        entry.get("exploitedSince")
        or entry.get("exploited_since")
        or entry.get("datePublished")
        or entry.get("date_published")
        or entry.get("dateAdded")
        or entry.get("date_added")
        or entry.get("published")
    )
    date_added = _parse_euvd_date(raw_date) if raw_date else None

    # Collect CVE-bearing strings from common fields.
    cve_candidates: List[str] = []

    for key in ("aliases", "references", "description"):
        v = entry.get(key)
        if isinstance(v, str):
            cve_candidates.append(v)  # ENISA uses newline-separated strings
        elif isinstance(v, list):
            for a in v:
                if isinstance(a, str):
                    cve_candidates.append(a)
                elif isinstance(a, dict):
                    cve_candidates.extend(x for x in a.values() if isinstance(x, str))

    for key in ("cveId", "cve_id", "cve"):
        v = entry.get(key)
        if isinstance(v, str):
            cve_candidates.append(v)
        elif isinstance(v, list):
            cve_candidates.extend(x for x in v if isinstance(x, str))

    # Last-resort: regex over JSON dump (catches CVE strings buried anywhere)
    if not cve_candidates:
        cve_candidates.append(json.dumps(entry, default=str))

    cves = set()
    for s in cve_candidates:
        for match in CVE_REGEX.findall(s):
            cves.add(match.upper())

    for cve in cves:
        yield cve, {
            "date_added": date_added,
            "critical": is_critical,
            "euvd_id": euvd_id,
        }


def _extract_entries(payload) -> List[dict]:
    """Find the list-of-entries inside an ENISA response (varying shapes)."""
    if isinstance(payload, list):
        return [e for e in payload if isinstance(e, dict)]
    if isinstance(payload, dict):
        for key in ("items", "data", "vulnerabilities", "results"):
            v = payload.get(key)
            if isinstance(v, list):
                return [e for e in v if isinstance(e, dict)]
        if "id" in payload:
            return [payload]
    return []


def _merge_euvd_entries(
    entries: List[dict],
    is_critical: bool,
    merged: Dict[str, dict],
) -> int:
    """Apply parser to entries, merge into `merged` map. Returns count of entries processed."""
    count = 0
    for entry in entries:
        for cve, info in _parse_euvd_entry(entry, is_critical):
            existing = merged.get(cve)
            if existing:
                existing["critical"] = existing.get("critical") or info["critical"]
                if not existing.get("date_added") and info.get("date_added"):
                    existing["date_added"] = info["date_added"]
                if not existing.get("euvd_id") and info.get("euvd_id"):
                    existing["euvd_id"] = info["euvd_id"]
            else:
                merged[cve] = info
        count += 1
    return count


def _fetch_euvd_search_paginated(
    client: httpx.Client,
    base_url: str,
    is_critical: bool,
    merged: Dict[str, dict],
) -> int:
    """
    Fetch all pages from a /search?... endpoint.
    Returns total entries processed.
    """
    total_processed = 0
    sep = "&" if "?" in base_url else "?"
    for page in range(EUVD_MAX_PAGES):
        page_url = f"{base_url}{sep}size={EUVD_PAGE_SIZE}&page={page}"
        try:
            resp = client.get(page_url, headers={"Accept": "application/json"})
            resp.raise_for_status()
            payload = resp.json()
        except httpx.HTTPError as e:
            logger.warning(f"EUVD search page {page} failed: {e}")
            break

        entries = _extract_entries(payload)
        if not entries:
            break

        processed = _merge_euvd_entries(entries, is_critical, merged)
        total_processed += processed

        # Stop when we've covered the reported total
        if isinstance(payload, dict):
            reported_total = payload.get("total")
            if isinstance(reported_total, int) and (page + 1) * EUVD_PAGE_SIZE >= reported_total:
                break

        if len(entries) < EUVD_PAGE_SIZE:
            break  # last page

    return total_processed


def fetch_euvd_catalogs(db: Session, force_refresh: bool = False) -> Dict[str, dict]:
    """
    Fetch ENISA EUVD catalogs and merge into one CVE map.
    Cached for 24h in `settings` table.

    Sources:
    - /api/search?exploited=true  — paginated, ~1500+ CVEs (CISA-KEV-style)
    - /api/criticalvulnerabilities — small curated ENISA-priority list (Critical flag)

    Returns: { "CVE-XXXX-NNNN": {date_added, critical, euvd_id}, ... }
    """
    if not force_refresh:
        cached = _load_euvd_cache(db)
        if cached is not None:
            logger.debug(f"EUVD: using cached catalog ({len(cached)} entries)")
            return cached

    logger.info("EUVD: fetching fresh catalogs from ENISA")
    merged: Dict[str, dict] = {}
    exploited_count = 0
    critical_count = 0

    try:
        with httpx.Client(timeout=HTTP_TIMEOUT) as client:
            # 1) Exploited (paginated)
            exploited_count = _fetch_euvd_search_paginated(
                client, EUVD_SEARCH_EXPLOITED_URL, is_critical=False, merged=merged
            )

            # 2) Curated Critical (small list, single call)
            try:
                resp = client.get(EUVD_CRITICAL_URL, headers={"Accept": "application/json"})
                resp.raise_for_status()
                critical_count = _merge_euvd_entries(
                    _extract_entries(resp.json()), is_critical=True, merged=merged
                )
            except httpx.HTTPError as e:
                logger.warning(f"EUVD critical fetch failed: {e}")
    except Exception as e:
        logger.error(f"EUVD client error: {e}")

    if not merged:
        # All fetches failed: try stale cache
        cache_setting = db.query(Setting).filter(Setting.key == EUVD_CACHE_KEY).first()
        if cache_setting and cache_setting.value:
            try:
                logger.warning("EUVD: returning stale cache after fetch failure")
                return json.loads(cache_setting.value)
            except json.JSONDecodeError:
                pass
        raise EnrichmentError("EUVD fetch failed and no cache available")

    _store_euvd_cache(db, merged)
    logger.info(
        f"EUVD: cached {len(merged)} unique CVEs from ENISA "
        f"(exploited pages processed={exploited_count}, critical entries={critical_count})"
    )
    return merged


# ============================================================
# Apply to DB
# ============================================================

def enrich_vulnerabilities(
    db: Session,
    vulns: List[Vulnerability],
    use_epss: Optional[bool] = None,
    use_kev: Optional[bool] = None,
    use_euvd: Optional[bool] = None,
) -> dict:
    """
    Reichert eine Liste von Vulnerabilities in-place an und commited.

    Returns stats dict: {epss_updated, kev_marked, kev_cleared,
                         euvd_marked, euvd_cleared, total}.
    """
    if not vulns:
        return {
            "epss_updated": 0,
            "kev_marked": 0, "kev_cleared": 0,
            "euvd_marked": 0, "euvd_cleared": 0,
            "total": 0,
        }

    if use_epss is None:
        use_epss = _setting_bool(db, SETTING_EPSS_ENABLED, default=True)
    if use_kev is None:
        use_kev = _setting_bool(db, SETTING_KEV_ENABLED, default=True)
    if use_euvd is None:
        use_euvd = _setting_bool(db, SETTING_EUVD_ENABLED, default=True)

    stats = {
        "epss_updated": 0,
        "kev_marked": 0, "kev_cleared": 0,
        "euvd_marked": 0, "euvd_cleared": 0,
        "total": len(vulns),
    }

    # Deduplicate CVE IDs (multiple assets can share a CVE)
    cve_ids = sorted({v.cve_id for v in vulns if v.cve_id})

    epss_map: Dict[str, Dict[str, float]] = {}
    if use_epss and cve_ids:
        try:
            epss_map = fetch_epss_scores(cve_ids)
        except Exception as e:
            logger.error(f"EPSS enrichment failed: {e}")

    kev_map: Dict[str, dict] = {}
    if use_kev:
        try:
            kev_map = fetch_kev_catalog(db)
        except EnrichmentError as e:
            logger.error(f"KEV enrichment skipped: {e}")
        except Exception as e:
            logger.error(f"KEV enrichment failed unexpectedly: {e}")

    euvd_map: Dict[str, dict] = {}
    if use_euvd:
        try:
            euvd_map = fetch_euvd_catalogs(db)
        except EnrichmentError as e:
            logger.error(f"EUVD enrichment skipped: {e}")
        except Exception as e:
            logger.error(f"EUVD enrichment failed unexpectedly: {e}")

    now = datetime.now()

    for vuln in vulns:
        sources_used = []

        if use_epss and vuln.cve_id in epss_map:
            data = epss_map[vuln.cve_id]
            vuln.epss_score = data["epss"]
            vuln.epss_percentile = data.get("percentile")
            vuln.epss_updated_at = now
            sources_used.append("epss")
            stats["epss_updated"] += 1

        if use_kev:
            if vuln.cve_id in kev_map:
                kev_entry = kev_map[vuln.cve_id]
                if not vuln.kev_listed:
                    stats["kev_marked"] += 1
                vuln.kev_listed = True
                vuln.kev_ransomware_use = bool(kev_entry.get("ransomware_use"))
                vuln.kev_short_description = kev_entry.get("short_description")
                date_added_str = kev_entry.get("date_added")
                if date_added_str:
                    try:
                        vuln.kev_date_added = datetime.fromisoformat(date_added_str)
                    except ValueError:
                        pass
                sources_used.append("kev")
            else:
                if vuln.kev_listed:
                    # No longer in KEV (rare but possible)
                    vuln.kev_listed = False
                    vuln.kev_ransomware_use = False
                    vuln.kev_date_added = None
                    vuln.kev_short_description = None
                    stats["kev_cleared"] += 1

        if use_euvd:
            if vuln.cve_id in euvd_map:
                euvd_entry = euvd_map[vuln.cve_id]
                if not vuln.euvd_listed:
                    stats["euvd_marked"] += 1
                vuln.euvd_listed = True
                vuln.euvd_critical = bool(euvd_entry.get("critical"))
                vuln.euvd_id = euvd_entry.get("euvd_id")
                date_added_str = euvd_entry.get("date_added")
                if date_added_str:
                    # ENISA dates can be plain dates or ISO timestamps
                    try:
                        vuln.euvd_date_added = datetime.fromisoformat(
                            date_added_str.replace("Z", "+00:00")
                        )
                    except (ValueError, AttributeError):
                        pass
                sources_used.append("euvd")
            else:
                if vuln.euvd_listed:
                    vuln.euvd_listed = False
                    vuln.euvd_critical = False
                    vuln.euvd_date_added = None
                    vuln.euvd_id = None
                    stats["euvd_cleared"] += 1

        if sources_used:
            # Merge with existing sources
            existing = []
            if vuln.enrichment_sources:
                try:
                    existing = json.loads(vuln.enrichment_sources)
                except json.JSONDecodeError:
                    existing = []
            vuln.enrichment_sources = json.dumps(sorted(set(existing) | set(sources_used)))
            vuln.enrichment_updated_at = now

    db.commit()
    logger.info(
        f"Enrichment done: {stats['total']} vulns, "
        f"epss_updated={stats['epss_updated']}, "
        f"kev_marked={stats['kev_marked']}, kev_cleared={stats['kev_cleared']}, "
        f"euvd_marked={stats['euvd_marked']}, euvd_cleared={stats['euvd_cleared']}"
    )
    return stats


def enrich_vulnerability_by_id(db: Session, vuln_id: int) -> dict:
    """Single-vuln convenience wrapper."""
    vuln = db.query(Vulnerability).filter(Vulnerability.id == vuln_id).first()
    if not vuln:
        raise EnrichmentError(f"Vulnerability {vuln_id} not found")
    return enrich_vulnerabilities(db, [vuln])


def enrich_all_open_vulnerabilities(db: Session) -> dict:
    """Refresh enrichment for every open vuln. Use for scheduled daily refresh."""
    from app.models.vulnerability import VulnerabilityStatus

    vulns = db.query(Vulnerability).filter(
        Vulnerability.status == VulnerabilityStatus.open
    ).all()
    return enrich_vulnerabilities(db, vulns)