vuln-nvd/scripts/sync.py

#!/usr/bin/env python3
"""NVD CVE API 2.0 mirror. Requires Python >=3.10, no third-party modules."""
import argparse
import json
import os
import random
import re
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timedelta, timezone
from pathlib import Path

API = 'https://services.nvd.nist.gov/rest/json/cves/2.0'
ROOT = Path(__file__).resolve().parents[1]
CVE_RE = re.compile(r'^CVE-(\d{4})-\d{4,}$')
UTC = timezone.utc

def timestamp(dt):
    return dt.astimezone(UTC).strftime('%Y-%m-%dT%H:%M:%S.000+00:00')

def parse_time(s):
    return datetime.fromisoformat(s.replace('Z', '+00:00')).astimezone(UTC)

def atomic_json(path, obj):
    path.parent.mkdir(parents=True, exist_ok=True)
    content = (json.dumps(obj, ensure_ascii=False, separators=(',', ':'), sort_keys=True) + '\n').encode('utf-8')
    tmp = path.with_name(path.name + '.tmp')
    with tmp.open('wb') as f:
        f.write(content)
    os.replace(tmp, path)

def windows(start, end, days=30):
    while start < end:
        nxt = min(start + timedelta(days=days), end)
        yield start, nxt
        start = nxt

class Client:
    def __init__(self, api_key=None, delay=None, retries=6):
        self.api_key = api_key
        self.delay = delay if delay is not None else (0.8 if api_key else 6.5)
        self.retries = retries
        self.last_call = 0

    def fetch(self, params):
        query = urllib.parse.urlencode(params)
        url = API + '?' + query
        headers = {'User-Agent': 'vuln-nvd-mirror/1.0 (+https://git.apex-project.de/VULN/vuln-nvd)', 'Accept': 'application/json'}
        if self.api_key:
            headers['apiKey'] = self.api_key
        for attempt in range(self.retries + 1):
            wait = self.delay - (time.monotonic() - self.last_call)
            if wait > 0:
                time.sleep(wait)
            self.last_call = time.monotonic()
            try:
                with urllib.request.urlopen(urllib.request.Request(url, headers=headers), timeout=120) as resp:
                    data = json.load(resp)
                if not isinstance(data, dict) or not isinstance(data.get('vulnerabilities'), list) or not isinstance(data.get('totalResults'), int):
                    raise ValueError('Unexpected NVD API JSON payload')
                return data
            except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError) as exc:
                status = getattr(exc, 'code', None)
                if status is not None and status not in (403, 408, 429, 500, 502, 503, 504):
                    raise
                if attempt >= self.retries:
                    raise
                retry_after = None
                if isinstance(exc, urllib.error.HTTPError):
                    retry_after = exc.headers.get('Retry-After')
                try:
                    backoff = min(300, max(float(retry_after), 1)) if retry_after else min(240, 10 * 2**attempt + random.random() * 4)
                except ValueError:
                    backoff = min(240, 10 * 2**attempt)
                print(f'API error {exc}; retry {attempt+1}/{self.retries} after {backoff:.0f}s', file=sys.stderr, flush=True)
                time.sleep(backoff)

def import_period(client, start, end, date_prefix, min_year, counts):
    params = {date_prefix+'StartDate': timestamp(start), date_prefix+'EndDate': timestamp(end), 'resultsPerPage': 2000}
    offset = 0
    while True:
        response = client.fetch({**params, 'startIndex': offset})
        records = response['vulnerabilities']
        total = response['totalResults']
        if not records and offset < total:
            raise RuntimeError(f'Empty page before totalResults at offset {offset}/{total}')
        for item in records:
            if not isinstance(item, dict) or not isinstance(item.get('cve'), dict):
                raise ValueError('Malformed NVD vulnerability')
            cve = item['cve']
            identifier = cve.get('id', '')
            match = CVE_RE.fullmatch(identifier)
            if not match:
                raise ValueError(f'Invalid CVE ID: {identifier!r}')
            if int(match.group(1)) < min_year:
                counts['skipped'] += 1
                continue
            destination = ROOT / 'data' / match.group(1) / f'{identifier}.json'
            payload = (json.dumps(cve, sort_keys=True, ensure_ascii=False, separators=(',', ':')) + '\n').encode('utf-8')
            if destination.exists() and destination.read_bytes() == payload:
                counts['unchanged'] += 1
            else:
                atomic_json(destination, cve)
                counts['written'] += 1
        counts['seen'] += len(records)
        offset += len(records)
        print(f'  {date_prefix}: {start.date()} to {end.date()}, records {offset}/{total}', flush=True)
        if offset >= total:
            break
    return counts

def run(args):
    if args.since_year < 1999:
        raise ValueError('--since-year must be >= 1999')
    statefile = ROOT / 'state' / 'sync.json'
    state = json.loads(statefile.read_text(encoding='utf-8')) if statefile.exists() else {}
    end = parse_time(args.end) if args.end else datetime.now(UTC) - timedelta(minutes=5)
    end = end.replace(microsecond=0)
    client = Client(api_key=os.getenv('NVD_API_KEY'), delay=args.delay)
    mode = args.mode if args.mode != 'auto' else ('incremental' if state.get('initial_complete') else 'full')
    if mode == 'incremental' and not state.get('initial_complete'):
        raise ValueError('Initial full import required before incremental mode')
    start = datetime(args.since_year, 1, 1, tzinfo=UTC) if mode == 'full' else parse_time(state['last_successful_end']) - timedelta(hours=args.overlap_hours)
    if start >= end:
        print('Already up to date')
        return
    counts = {'seen': 0, 'written': 0, 'unchanged': 0, 'skipped': 0}
    prefix = 'pub' if mode == 'full' else 'lastMod'
    print(f'Mode={mode} from={timestamp(start)} to={timestamp(end)} with API key={bool(client.api_key)}', flush=True)
    for a, b in windows(start, end, args.window_days):
        import_period(client, a, b, prefix, args.since_year, counts)
    atomic_json(statefile, {
        'schema_version': 1,
        'source': 'NVD CVE API 2.0',
        'since_year': args.since_year,
        'initial_complete': True,
        'last_successful_end': end.isoformat(),
        'last_run_mode': mode,
        'last_run_counts': counts,
        'updated_at': datetime.now(UTC).isoformat()
    })
    print('Success:', json.dumps(counts), flush=True)

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--since-year', type=int, default=2017)
    parser.add_argument('--mode', choices=('auto', 'full', 'incremental'), default='auto')
    parser.add_argument('--overlap-hours', type=int, default=48)
    parser.add_argument('--window-days', type=int, default=30)
    parser.add_argument('--delay', type=float, default=None)
    parser.add_argument('--end', help='Optional UTC ISO timestamp (testing/backfill)')
    args = parser.parse_args()
    if not 1 <= args.window_days <= 119 or args.overlap_hours < 0:
        parser.error('window-days must be 1..119; overlap-hours must be nonnegative')
    try:
        run(args)
    except Exception as exc:
        print(f'FAILED, sync state not advanced: {exc}', file=sys.stderr)
        sys.exit(1)

if __name__ == '__main__':
    main()