#!/usr/bin/env python3
"""Incrementally mirror CVEProject/cvelistV5 CVE records into data/.

Requires Linux, git >=2.25, rsync, Python 3.9+.
Persistent upstream clone is protected by flock; no new clone after first run.
"""

import argparse
import datetime as dt
import fcntl
import json
import os
from pathlib import Path
import re
import subprocess

ROOT = Path(__file__).resolve().parents[1]
UPSTREAM_URL = 'https://github.com/CVEProject/cvelistV5.git'
DEFAULT_CACHE = Path('/var/cache/vuln-cve')


def command(*args, cwd=None, capture=False):
    return subprocess.run(
        list(map(str, args)), cwd=cwd, check=True,
        text=True, capture_output=capture,
    )


def git(cwd, *args):
    return command('git', '-C', cwd, *args, capture=True).stdout.strip()


def prepare_cache(path, source_url, years):
    if not (path / '.git').is_dir():
        if path.exists() and any(path.iterdir()):
            raise RuntimeError(f'Incomplete cache directory: {path}; move it aside and retry')
        path.parent.mkdir(parents=True, exist_ok=True)
        command('git', 'clone', '--depth=1', '--filter=blob:none', '--sparse',
                '--branch=main', source_url, path)
        command('git', '-C', path, 'sparse-checkout', 'set', *years)
    else:
        # Refuse accidentally reusing a different upstream repository.
        remote_url = git(path, 'remote', 'get-url', 'origin')
        if remote_url != source_url:
            raise RuntimeError(f'Cache remote mismatch: {remote_url}')
        existing = git(path, 'sparse-checkout', 'list').splitlines()
        if existing != years:
            command('git', '-C', path, 'sparse-checkout', 'set', *years)

    # Shallow fetch downloads objects missing from persistent cache.
    command('git', '-C', path, 'fetch', '--depth=1', '--no-tags', 'origin', 'main')
    command('git', '-C', path, 'reset', '--hard', 'FETCH_HEAD')
    return git(path, 'rev-parse', 'HEAD')


def validate_records(source_root, minimum):
    count = 0
    for file in source_root.rglob('CVE-*.json'):
        if not re.fullmatch(r'CVE-\d{4}-\d{4,}', file.stem):
            raise RuntimeError(f'Invalid filename: {file}')
        try:
            with file.open('r', encoding='utf-8') as handle:
                doc = json.load(handle)
        except (ValueError, OSError) as exc:
            raise RuntimeError(f'Invalid JSON file {file}: {exc}') from exc
        if doc.get('cveMetadata', {}).get('cveId') != file.stem:
            raise RuntimeError(f'CVE ID mismatch: {file}')
        count += 1
    if count < minimum:
        raise RuntimeError(f'Only {count} records (< {minimum}): refusing destructive sync')
    return count


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--since-year', type=int, default=2017)
    parser.add_argument('--cache-dir', type=Path, default=Path(os.environ.get('CVE_CACHE_DIR', DEFAULT_CACHE)))
    parser.add_argument('--source-url', default=UPSTREAM_URL)
    parser.add_argument('--minimum-records', type=int, default=1000)
    parser.add_argument('--force-verify', action='store_true', help='Validate and rsync even if upstream commit has not changed')
    args = parser.parse_args()

    current_year = dt.datetime.now(dt.timezone.utc).year
    if not 1999 <= args.since_year <= current_year:
        parser.error('Invalid --since-year')
    years = [f'cves/{year}' for year in range(args.since_year, current_year + 1)]

    cache_root = args.cache_dir.resolve()
    cache_root.mkdir(parents=True, exist_ok=True)
    with (cache_root / 'sync.lock').open('a+') as lock:
        fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
        upstream = cache_root / 'upstream'
        revision = prepare_cache(upstream, args.source_url, years)

        data_dir = ROOT / 'data'
        state_file = ROOT / 'state' / 'sync.json'
        previous = {}
        if state_file.is_file():
            previous = json.loads(state_file.read_text(encoding='utf-8'))
        if (
            not args.force_verify
            and previous.get('upstream_commit') == revision
            and previous.get('since_year') == args.since_year
            and data_dir.is_dir()
            and any(data_dir.rglob('CVE-*.json'))
        ):
            print(f'No upstream changes; cached revision {revision}')
            return

        src = upstream / 'cves'
        records = validate_records(src, args.minimum_records)
        data_dir.mkdir(parents=True, exist_ok=True)
        command('rsync', '-a', '--delete', '--exclude=.gitkeep', f'{src}/', f'{data_dir}/')

        state_file.parent.mkdir(parents=True, exist_ok=True)
        state = {
            'source': args.source_url,
            'upstream_commit': revision,
            'since_year': args.since_year,
            'records': records,
            'synced_at_utc': dt.datetime.now(dt.timezone.utc).isoformat(),
        }
        tmp = state_file.with_suffix('.json.tmp')
        tmp.write_text(json.dumps(state, indent=2, ensure_ascii=False) + '\n', encoding='utf-8')
        os.replace(tmp, state_file)
        print(f'Mirrored {records} records from {revision}')


if __name__ == '__main__':
    main()