#!/usr/bin/env python3
"""
linux_updates_threatdb - threat data for the linux_updates check plug-ins

Author:   Marius Gielnik
Project:  https://github.com/Majus420/checkmk/tree/main/MKP/advanced_linux_update_check
Version:  1.3.1
License:  GPL-2.0-or-later

Purpose:  Downloads vulnerability and lifecycle data for the check
          "<PM> Pending Updates" / "<PM> Repository Health":
            - cves.json:   CISA KEV list + EPSS scores >= 1 %   (all distributions,
                           matched against the CVEs of dnf/yum/zypper advisories)
            - <release>.json: CVE -> fixed version per source package from OSV
                           (Debian/Ubuntu, matched by version comparison)
            - eol.json:    end-of-life dates from endoflife.date
Sources:  CISA KEV       https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json
          EPSS (FIRST)   https://epss.empiricalsecurity.com/epss_scores-current.csv.gz
          OSV Debian     https://osv-vulnerabilities.storage.googleapis.com/Debian/all.zip
          OSV Ubuntu     https://osv-vulnerabilities.storage.googleapis.com/Ubuntu/all.zip
          endoflife.date https://endoflife.date/api/<product>.json
          Downloads are conditional (ETag/Last-Modified), a failed download
          keeps the last good copy; the age of each source is reported by
          the service "APT Repository Health".
Runs as:  site user, on EVERY site that monitors Debian/Ubuntu hosts
Schedule: no cron needed. The check plug-ins start this script automatically
          in the background on every site that monitors Linux hosts, as soon
          as the threat data is older than 6 hours (at most one attempt per
          hour, a lock prevents parallel runs). A new site therefore gets its
          data without any setup. The first run takes a few minutes.
Files:    ~/var/check_mk/linux_updates/threatdb/meta.json      status per source
          ~/var/check_mk/linux_updates/threatdb/cves.json      KEV + EPSS
          ~/var/check_mk/linux_updates/threatdb/eol.json       end-of-life dates
          ~/var/check_mk/linux_updates/threatdb/<release>.json e.g. debian_12, ubuntu_24.04
          ~/var/check_mk/linux_updates/download/               raw downloads (cache)
          ~/var/log/linux_updates_threatdb.log                 log (one line per run/error)
          ~/var/check_mk/linux_updates/trigger.stamp           last automatic start
Proxy:    sites without direct internet access: either the usual environment
          variables (HTTPS_PROXY / https_proxy, e.g. in ~/etc/environment) or
          ~/etc/linux_updates_threatdb.conf with the line
              proxy = http://proxy.example.com:3128
          (the config file wins over the environment)
Usage:    linux_updates_threatdb                 update now (manual, optional)
          linux_updates_threatdb --status        short status
Removal:  remove the extension package; rm -r ~/var/check_mk/linux_updates
"""

from __future__ import annotations

import argparse
import csv
import fcntl
import gzip
import json
import os
import re
import sys
import time
import urllib.error
import urllib.request
import zipfile
from pathlib import Path
from typing import Any

OMD_ROOT = Path(os.environ.get("OMD_ROOT", ""))
if not str(OMD_ROOT) or not OMD_ROOT.is_dir():
    sys.exit("linux_updates_threatdb: must run as site user (OMD_ROOT not set)")

sys.path.insert(0, str(OMD_ROOT / "local/lib/python3"))
from cmk_addons.plugins.linux_updates.lib.common import (  # noqa: E402
    EOL_PRODUCTS,
    META_FILE,
    THREATDB_DIR,
    ecosystem_from_osv,
)

VERSION = "1.3.1"
BASE = OMD_ROOT / "var/check_mk/linux_updates"
DOWNLOAD = BASE / "download"
LOG = OMD_ROOT / "var/log/linux_updates_threatdb.log"
LOCK = BASE / "update.lock"
CONF = OMD_ROOT / "etc/linux_updates_threatdb.conf"

SOURCES = {
    "kev": "https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json",
    "epss": "https://epss.empiricalsecurity.com/epss_scores-current.csv.gz",
    "osv_debian": "https://osv-vulnerabilities.storage.googleapis.com/Debian/all.zip",
    "osv_ubuntu": "https://osv-vulnerabilities.storage.googleapis.com/Ubuntu/all.zip",
}
FILENAMES = {
    "kev": "kev.json",
    "epss": "epss.csv.gz",
    "osv_debian": "osv_debian.zip",
    "osv_ubuntu": "osv_ubuntu.zip",
}
# entries kept in the condensed data (everything else is plain "security")
KEEP_SEVERITIES = {"critical", "high", "medium"}
EPSS_FLOOR = 0.01
CVE_RE = re.compile(r"CVE-\d{4}-\d{4,}")
TIMEOUT = 120
EOL_URL = "https://endoflife.date/api/{product}.json"
EOL_LIST = sorted(set(EOL_PRODUCTS.values()) | {"centos", "centos-stream"})


def setup_proxy() -> str:
    """Proxy from ~/etc/linux_updates_threatdb.conf, else environment (urllib default)."""
    proxy = ""
    try:
        for line in CONF.read_text(encoding="utf-8").splitlines():
            key, _, value = line.partition("=")
            if key.strip() == "proxy" and value.strip():
                proxy = value.strip()
    except OSError:
        pass
    if proxy:
        handler = urllib.request.ProxyHandler({"http": proxy, "https": proxy})
        urllib.request.install_opener(urllib.request.build_opener(handler))
        return proxy
    env = urllib.request.getproxies()
    return env.get("https", env.get("http", ""))


def log(msg: str) -> None:
    line = f"{time.strftime('%Y-%m-%d %H:%M:%S')} {msg}"
    try:
        LOG.parent.mkdir(parents=True, exist_ok=True)
        with LOG.open("a", encoding="utf-8") as fh:
            fh.write(line + "\n")
    except OSError:
        pass
    print(line)


def write_atomic(path: Path, data: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_bytes(data)
    tmp.replace(path)


def load_json(path: Path, default: Any) -> Any:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return default


# ---------------------------------------------------------------- download


def download(name: str, meta: dict[str, Any]) -> bool:
    """Conditional download into DOWNLOAD; True if a usable local copy exists."""
    url = SOURCES[name]
    target = DOWNLOAD / FILENAMES[name]
    state = meta.setdefault("sources", {}).setdefault(name, {})
    state["last_attempt"] = time.time()
    req = urllib.request.Request(url, headers={"User-Agent": f"linux_updates_threatdb/{VERSION}"})
    if target.exists():
        if state.get("etag"):
            req.add_header("If-None-Match", state["etag"])
        if state.get("last_modified"):
            req.add_header("If-Modified-Since", state["last_modified"])
    try:
        with urllib.request.urlopen(req, timeout=TIMEOUT) as resp:
            tmp = target.with_suffix(target.suffix + ".part")
            target.parent.mkdir(parents=True, exist_ok=True)
            with tmp.open("wb") as fh:
                while chunk := resp.read(1 << 20):
                    fh.write(chunk)
            tmp.replace(target)
            state["etag"] = resp.headers.get("ETag", "")
            state["last_modified"] = resp.headers.get("Last-Modified", "")
        state["last_success"] = time.time()
        state["error"] = ""
        return True
    except urllib.error.HTTPError as exc:
        if exc.code == 304 and target.exists():
            state["last_success"] = time.time()
            state["error"] = ""
            return True
        state["error"] = f"HTTP {exc.code}"
    except (urllib.error.URLError, OSError, TimeoutError) as exc:
        state["error"] = str(exc)[:200]
    log(f"download {name} failed: {state['error']} (using cached copy: {target.exists()})")
    return target.exists()


# ------------------------------------------------------------------- parse


def read_kev() -> set[str]:
    data = load_json(DOWNLOAD / FILENAMES["kev"], {})
    return {v.get("cveID", "") for v in data.get("vulnerabilities", []) if v.get("cveID")}


def read_epss() -> dict[str, float]:
    result: dict[str, float] = {}
    try:
        with gzip.open(DOWNLOAD / FILENAMES["epss"], "rt", encoding="utf-8") as fh:
            lines = (line for line in fh if not line.startswith("#"))
            for row in csv.DictReader(lines):
                try:
                    result[row["cve"]] = float(row["epss"])
                except (KeyError, ValueError):
                    continue
    except (OSError, EOFError):
        pass
    return result


def _severity(record: dict[str, Any]) -> str:
    for sev in record.get("severity") or []:
        if str(sev.get("type", "")).lower() in ("ubuntu", "debian"):
            return str(sev.get("score", "")).lower()
    dbs = record.get("database_specific") or {}
    for key in ("priority", "urgency", "severity"):
        if isinstance(dbs.get(key), str):
            return dbs[key].lower()
    return ""


def _cves(record: dict[str, Any]) -> set[str]:
    found: set[str] = set()
    for value in [record.get("id", "")] + list(record.get("aliases") or []) + list(record.get("upstream") or []):
        found.update(CVE_RE.findall(str(value)))
    return found


def read_osv(name: str, kev: set[str], epss: dict[str, float], out: dict[str, dict[str, set]]) -> int:
    count = 0
    path = DOWNLOAD / FILENAMES[name]
    try:
        zf = zipfile.ZipFile(path)
    except (OSError, zipfile.BadZipFile) as exc:
        log(f"cannot read {path.name}: {exc}")
        return 0
    with zf:
        for info in zf.infolist():
            if not info.filename.endswith(".json"):
                continue
            try:
                record = json.loads(zf.read(info))
            except (ValueError, KeyError, zipfile.BadZipFile):
                continue
            if record.get("withdrawn"):
                continue
            cves = _cves(record)
            if not cves:
                continue
            sev = _severity(record)
            for aff in record.get("affected") or []:
                pkg = aff.get("package") or {}
                eco = ecosystem_from_osv(str(pkg.get("ecosystem", "")))
                src = pkg.get("name")
                if not eco or not src:
                    continue
                for rng in aff.get("ranges") or []:
                    for event in rng.get("events") or []:
                        fixed = event.get("fixed")
                        if not fixed:
                            continue
                        for cve in cves:
                            is_kev = cve in kev
                            score = epss.get(cve, 0.0)
                            if not (is_kev or sev in KEEP_SEVERITIES or score >= EPSS_FLOOR):
                                continue
                            out.setdefault(eco, {}).setdefault(src, set()).add(
                                (fixed, cve, 1 if is_kev else 0, sev, round(score, 4))
                            )
                            count += 1
    return count


# --------------------------------------------------------------------- eol


def update_eol(meta: dict[str, Any]) -> None:
    """endoflife.date -> eol.json {"products": {product: {cycle: "YYYY-MM-DD"}}}"""
    state = meta.setdefault("sources", {}).setdefault("eol", {})
    state["last_attempt"] = time.time()
    old = load_json(THREATDB_DIR / "eol.json", {}).get("products", {})
    products: dict[str, dict[str, str]] = {}
    errors = []
    for product in EOL_LIST:
        req = urllib.request.Request(
            EOL_URL.format(product=product), headers={"User-Agent": f"linux_updates_threatdb/{VERSION}"}
        )
        try:
            with urllib.request.urlopen(req, timeout=TIMEOUT) as resp:
                cycles = json.loads(resp.read())
            products[product] = {
                str(c.get("cycle")): str(c.get("eol"))
                for c in cycles
                if isinstance(c, dict) and isinstance(c.get("eol"), str)
            }
        except (urllib.error.URLError, OSError, TimeoutError, ValueError) as exc:
            errors.append(f"{product}: {str(exc)[:80]}")
            if product in old:
                products[product] = old[product]
    write_atomic(THREATDB_DIR / "eol.json", json.dumps({"generated": time.time(), "products": products}).encode())
    if len(errors) < len(EOL_LIST):
        state["last_success"] = time.time()
    state["error"] = "; ".join(errors)[:300]
    if errors:
        log(f"eol data incomplete: {state['error']}")


# ------------------------------------------------------------------ update


def update() -> int:
    BASE.mkdir(parents=True, exist_ok=True)
    THREATDB_DIR.mkdir(parents=True, exist_ok=True)
    with LOCK.open("w") as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except OSError:
            print("another update is running")
            return 0
        meta = load_json(THREATDB_DIR / META_FILE, {})
        started = time.time()
        available = {name: download(name, meta) for name in SOURCES}
        kev = read_kev() if available["kev"] else set()
        epss = read_epss() if available["epss"] else {}

        out: dict[str, dict[str, set]] = {}
        entries = 0
        for name in ("osv_debian", "osv_ubuntu"):
            if available[name]:
                entries += read_osv(name, kev, epss, out)
        write_atomic(
            THREATDB_DIR / "cves.json",
            json.dumps(
                {
                    "generated": time.time(),
                    "kev": sorted(kev),
                    "epss": {k: round(v, 4) for k, v in epss.items() if v >= EPSS_FLOOR},
                },
                separators=(",", ":"),
            ).encode(),
        )
        update_eol(meta)

        if not out:
            meta["generated"] = meta.get("generated", 0)
            write_atomic(THREATDB_DIR / META_FILE, json.dumps(meta, indent=1).encode())
            log("no vulnerability data available, previous threat data kept")
            return 1

        for eco, pkgs in out.items():
            data = {
                "generated": time.time(),
                "packages": {src: sorted(map(list, items)) for src, items in sorted(pkgs.items())},
            }
            write_atomic(THREATDB_DIR / f"{eco}.json", json.dumps(data, separators=(",", ":")).encode())
        meta["generated"] = time.time()
        meta["version"] = VERSION
        meta["ecosystems"] = sorted(out)
        meta["counts"] = {"kev": len(kev), "epss": len(epss), "entries": entries}
        write_atomic(THREATDB_DIR / META_FILE, json.dumps(meta, indent=1).encode())
        errors = [f"{n}: {s.get('error')}" for n, s in meta.get("sources", {}).items() if s.get("error")]
        log(
            f"updated in {time.time() - started:.0f}s: {len(out)} releases, kev={len(kev)}, "
            f"epss={len(epss)}, entries={entries}" + (f", errors: {'; '.join(errors)}" if errors else "")
        )
        return 1 if errors else 0


def status() -> int:
    meta = load_json(THREATDB_DIR / META_FILE, None)
    if not meta:
        print("no threat data yet - run: linux_updates_threatdb")
        return 1
    now = time.time()
    gen = meta.get("generated", 0)
    print(f"generated: {time.strftime('%Y-%m-%d %H:%M', time.localtime(gen))} ({(now - gen) / 3600:.1f} h ago)")
    for name, s in sorted(meta.get("sources", {}).items()):
        age = (now - s.get("last_success", 0)) / 3600 if s.get("last_success") else -1
        print(f"  {name:11} last success {age:6.1f} h ago  {('ERROR ' + s['error']) if s.get('error') else 'ok'}")
    print(f"releases: {', '.join(meta.get('ecosystems', []))}")
    print(f"counts: {meta.get('counts', {})}")
    return 0




def main() -> int:
    parser = argparse.ArgumentParser(description="threat data for linux_updates")
    parser.add_argument("--status", action="store_true", help="show status")
    args = parser.parse_args()
    proxy = setup_proxy()
    if args.status:
        print(f"proxy: {proxy or 'none'}")
    if args.status:
        return status()
    return update()


if __name__ == "__main__":
    sys.exit(main())
