146 lines
6.4 KiB
Python
146 lines
6.4 KiB
Python
#!/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()
|