CVEs-PoC/scripts/build_site.py

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Dict, Tuple

from jinja2 import Environment, FileSystemLoader, select_autoescape

from utils import (
    API_DIR,
    DOCS_DIR,
    TEMPLATES_DIR,
    ensure_dirs,
    load_json,
    load_poc_index,
    parse_trending_from_readme,
    save_json,
)

from build_joined import build_joined, write_api_outputs
from build_diffs import build_diff, prune_snapshots

KEV_DATA = DOCS_DIR.parent / "data" / "kev.json"
EPSS_DATA = DOCS_DIR.parent / "data" / "epss.json"
README_PATH = DOCS_DIR.parent / "README.md"


def build_env() -> Environment:
    loader = FileSystemLoader(str(TEMPLATES_DIR))
    env = Environment(loader=loader, autoescape=select_autoescape(["html", "xml"]))
    env.trim_blocks = True
    env.lstrip_blocks = True
    return env


def render(env: Environment, template_name: str, context: Dict, output_path: Path) -> None:
    html = env.get_template(template_name).render(**context)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(html, encoding="utf-8")


def load_joined() -> Dict:
    kev = load_json(KEV_DATA, default={})
    epss = load_json(EPSS_DATA, default={})
    poc_index = load_poc_index()
    payload = build_joined(kev, epss, poc_index)
    write_api_outputs(payload)
    return payload


def write_snapshot(joined: Dict) -> Path:
    snapshot_path = API_DIR / "snapshots" / f"{joined['generated']}.json"
    ensure_dirs(snapshot_path.parent)
    save_json(snapshot_path, joined)
    return snapshot_path


def select_trending(readme_rows: list[dict]) -> list[dict]:
    """Pick the first 20 entries from the newest year table in README."""
    if not readme_rows:
        return []

    def parse_year(row: dict) -> int | None:
        try:
            return int(row.get("year"))
        except (TypeError, ValueError):
            return None

    years = [yr for yr in (parse_year(row) for row in readme_rows) if yr is not None]
    if not years:
        return []

    latest_year = max(years)
    selected: list[dict] = []
    for row in readme_rows:
        if parse_year(row) != latest_year:
            continue
        try:
            stars = int(row.get("stars") or 0)
        except (TypeError, ValueError):
            stars = 0
        selected.append(
            {
                "stars": stars,
                "updated": (row.get("updated") or "").strip(),
                "name": (row.get("name") or "").strip(),
                "url": (row.get("url") or "").strip(),
                "desc": (row.get("desc") or "").strip(),
                "year": latest_year,
            }
        )
        if len(selected) >= 20:
            break
    return selected


def build_pages(env: Environment, data: Dict, diff: Dict | None = None, html_mode: str = "summary") -> None:
    joined = data["joined"]
    details = data["details"]
    vendors = data["vendors"]

    trending_raw = parse_trending_from_readme(README_PATH)
    trending = select_trending(trending_raw)
    recent_kev = (diff or {}).get("new_kev_entries") or []
    metrics = {
        "kev_total": len(data["kev_enriched"]),
        "high_epss_count": len(joined["high_epss"]),
        "recent_kev_count": len(recent_kev),
    }

    if html_mode in {"summary", "all"}:
        common_ctx = {"generated": joined["generated"], "metrics": metrics, "recent_kev": recent_kev}
        render(
            env,
            "index.html",
            {**common_ctx, "data": joined, "trending": trending, "diff": diff or {}},
            DOCS_DIR / "index.html",
        )
        render(env, "kev.html", {**common_ctx, "kev": data["kev_enriched"]}, DOCS_DIR / "kev" / "index.html")
        render(env, "epss.html", {**common_ctx, "epss": joined["high_epss"]}, DOCS_DIR / "epss" / "index.html")
        render(env, "diffs.html", {**common_ctx, "diff": diff or {}}, DOCS_DIR / "diffs" / "index.html")

    if html_mode == "all":
        common_ctx = {"generated": joined["generated"]}
        for cve, detail in details.items():
            render(env, "cve.html", {**common_ctx, "cve": detail}, DOCS_DIR / "cve" / f"{cve}.html")

        for slug, vendor in vendors.items():
            cve_details = [details[cve] for cve in vendor["cves"] if cve in details]
            render(env, "vendor.html", {**common_ctx, "vendor": vendor, "cves": cve_details}, DOCS_DIR / "vendors" / f"{slug}.html")


def main() -> int:
    parser = argparse.ArgumentParser(description="Build static site and JSON")
    parser.add_argument(
        "--html-mode",
        choices=["none", "summary", "all"],
        default="none",
        help="Render no HTML, summary pages only, or all pages including per-CVE.",
    )
    args = parser.parse_args()

    ensure_dirs(DOCS_DIR, DOCS_DIR / "kev", DOCS_DIR / "epss", DOCS_DIR / "diffs")

    data = load_joined()
    # snapshot + diff before rendering so dashboard can show it
    snapshot_path = write_snapshot(data["joined"])
    snapshots = sorted((API_DIR / "snapshots").glob("*.json"))
    diff, target = build_diff(
        snapshots,
        kev_full=data["kev_enriched"],
        threshold=0.05,
        max_movers=50,
        recent_days=30,
    )
    prune_snapshots(snapshots, lookback_days=14)

    if args.html_mode != "none":
        env = build_env()
        build_pages(env, data, diff, html_mode=args.html_mode)

    # build daily diff after snapshot is written
    print("Site generated under docs/")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())