#!/usr/bin/env python3
"""Mirror raw NVD CPE dictionary and match criteria using resumable page checkpoints.
Requires Python 3.9+; no third-party dependencies.
"""
import argparse
import datetime as dt
import json
import os
from pathlib import Path
import random
import time
import urllib.error
import urllib.parse
import urllib.request

UTC = dt.timezone.utc
SOURCES = {
    "cpes": ("https://services.nvd.nist.gov/rest/json/cpes/2.0", "products", "cpe", "cpeNameId", 2000),
    "cpematch": ("https://services.nvd.nist.gov/rest/json/cpematch/2.0", "matchStrings", "matchString", "matchCriteriaId", 500),
}

def stamp(t):
    return t.astimezone(UTC).strftime("%Y-%m-%dT%H:%M:%S.000Z")

def parse_stamp(s):
    return dt.datetime.fromisoformat(s.replace("Z", "+00:00"))

def save_atomic(path, obj):
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = (json.dumps(obj, ensure_ascii=False, separators=(",", ":"), sort_keys=True) + "\n").encode("utf-8")
    if path.exists() and path.read_bytes() == payload:
        return False
    temp = path.with_suffix(path.suffix + ".tmp")
    temp.write_bytes(payload)
    os.replace(temp, path)
    return True

def fetch(url, params, api_key, delay, retries=5):
    query = urllib.parse.urlencode(params)
    request = urllib.request.Request(url + "?" + query, headers={
        "User-Agent": "vuln-cpe-mirror/1.0 (NVD API)", "Accept": "application/json",
        **({"apiKey": api_key} if api_key else {})
    })
    for attempt in range(retries):
        time.sleep(delay)
        try:
            with urllib.request.urlopen(request, timeout=120) as response:
                result = json.load(response)
            if not isinstance(result, dict):
                raise ValueError("NVD response not an object")
            return result
        except (urllib.error.URLError, TimeoutError, json.JSONDecodeError) as err:
            if isinstance(err, urllib.error.HTTPError) and err.code not in (403, 429, 500, 502, 503, 504):
                raise
            if attempt == retries - 1:
                raise
            wait = min(120, 2 ** (attempt + 2)) + random.random()
            print(f"Retry {attempt + 1}: {err}, sleep {wait:.1f}s", flush=True)
            time.sleep(wait)

def validate_page(page, key, nested, id_field):
    rows = page.get(key)
    if not isinstance(rows, list) or not isinstance(page.get("totalResults"), int):
        raise ValueError(f"Invalid NVD page envelope: {key}")
    if not isinstance(page.get("startIndex"), int):
        raise ValueError("Missing page startIndex")
    for wrapper in rows:
        if not isinstance(wrapper, dict) or not isinstance(wrapper.get(nested), dict):
            raise ValueError("Missing record object")
        obj = wrapper[nested]
        value = obj.get(id_field)
        if not isinstance(value, str) or not value:
            raise ValueError(f"Missing {id_field}")
    return rows

def process(name, root, api_key, max_pages, delay):
    endpoint, field, nested, id_field, page_size = SOURCES[name]
    state_path = root / "state" / f"{name}.json"
    state = json.loads(state_path.read_text()) if state_path.exists() else {"mode": "full", "offset": 0}
    now = dt.datetime.now(UTC)
    # First full sync runs across many jobs; delta begins at the *start* of the full scan.
    if state.get("mode") == "full" and not state.get("bootstrap_started"):
        state["bootstrap_started"] = stamp(now)
        save_atomic(state_path, state)
    processed = written = 0
    for _ in range(max_pages):
        if state["mode"] == "delta" and not state.get("window_end"):
            start = parse_stamp(state["last_completed"])
            end = min(start + dt.timedelta(days=110), dt.datetime.now(UTC) - dt.timedelta(minutes=2))
            if start >= end:
                break
            state["window_start"] = stamp(start)
            state["window_end"] = stamp(end)
            state["offset"] = 0
        params = {"startIndex": state.get("offset", 0), "resultsPerPage": page_size}
        if state["mode"] == "delta":
            params.update(lastModStartDate=state["window_start"], lastModEndDate=state["window_end"])
        result = fetch(endpoint, params, api_key, delay)
        offset = state.get("offset", 0)
        if result["startIndex"] != offset:
            raise RuntimeError(f"Unexpected pagination offset {result['startIndex']} instead of {offset}")
        rows = validate_page(result, field, nested, id_field)
        if not rows and offset < result["totalResults"]:
            raise RuntimeError("NVD returned an empty page before totalResults")
        for wrapper in rows:
            obj = wrapper[nested]
            rec_id = obj[id_field].lower()
            # NVD IDs are UUIDs; don't permit untrusted path separators.
            import uuid
            try: uid = str(uuid.UUID(rec_id))
            except ValueError as exc: raise ValueError(f"Bad {id_field}: {rec_id}") from exc
            path = root / "data" / name / uid[:2] / f"{uid}.json"
            written += save_atomic(path, wrapper)
            processed += 1
        next_offset = offset + len(rows)
        if next_offset >= result["totalResults"]:
            if state["mode"] == "full":
                state = {"mode": "delta", "last_completed": state["bootstrap_started"], "offset": 0}
            else:
                state = {"mode": "delta", "last_completed": state["window_end"], "offset": 0}
        else:
            state["offset"] = next_offset
        # State only moves forward after every returned record is durably written.
        save_atomic(state_path, state)
        print(f"{name}: +{len(rows)} records, written={written}, total={result['totalResults']}, phase={state['mode']}", flush=True)
        if state["mode"] == "delta" and not state.get("window_end"):
            # Only one delta window per run, never spin indefinitely.
            break
    return processed, written, state["mode"]

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=Path, default=Path("."))
    ap.add_argument("--max-pages", type=int, default=20)
    ap.add_argument("--delay", type=float)
    args = ap.parse_args()
    if not 1 <= args.max_pages <= 500:
        ap.error("--max-pages must be between 1 and 500")
    key = os.environ.get("NVD_API_KEY", "").strip()
    delay = args.delay if args.delay is not None else (0.7 if key else 6.2)
    for name in SOURCES:
        print(f"Start {name}", flush=True)
        print(process(name, args.root, key, args.max_pages, delay), flush=True)

if __name__ == "__main__": main()