#!/usr/bin/env python3
"""Mirror public BSI WID CSAF WHITE advisories using authoritative index.txt."""
import argparse
import hashlib
import json
import logging
import os
from pathlib import Path
import random
import re
import sys
import time
from datetime import datetime, timezone
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

BASE = 'https://wid.cert-bund.de/.well-known/csaf/white/'
PATH_RE = re.compile(r'^([12][0-9]{3})/([a-zA-Z0-9_.-]+\.json)$')
LOG = logging.getLogger('bsi')


def fetch(url, timeout=90, retries=5):
    for attempt in range(retries):
        try:
            req = Request(url, headers={'User-Agent': 'vuln-bsi-mirror/1.0 (+https://git.apex-project.de/VULN/vuln-bsi)', 'Accept': 'application/json,text/plain,*/*'})
            with urlopen(req, timeout=timeout) as response:
                return response.read()
        except (HTTPError, URLError, TimeoutError, OSError) as exc:
            if attempt + 1 == retries:
                raise RuntimeError(f'Failed to download {url}: {exc}') from exc
            delay = min(90, 2 ** attempt * 3 + random.random() * 2)
            LOG.warning('Download failed (%s), retry in %.1fs: %s', exc, delay, url)
            time.sleep(delay)


def parse_index(raw, min_year):
    try:
        lines = raw.decode('utf-8-sig').splitlines()
    except UnicodeError as exc:
        raise ValueError('index.txt is not valid UTF-8') from exc
    all_entries = [x.strip() for x in lines if x.strip()]
    if not all_entries:
        raise ValueError('Empty index; refusing to remove documents')
    if any(not PATH_RE.fullmatch(x) for x in all_entries):
        bad = next(x for x in all_entries if not PATH_RE.fullmatch(x))
        raise ValueError(f'Unexpected index entry: {bad[:140]!r}')
    if len(all_entries) != len(set(all_entries)):
        raise ValueError('Duplicate entries in index')
    selected = sorted(x for x in all_entries if int(x.split('/', 1)[0]) >= min_year)
    if not selected:
        raise ValueError('No entries in requested years')
    return selected


def validate_document(raw, path):
    try:
        obj = json.loads(raw)
    except (ValueError, UnicodeError) as exc:
        raise ValueError(f'Invalid JSON in {path}') from exc
    if not isinstance(obj, dict) or not isinstance(obj.get('document'), dict):
        raise ValueError(f'Not a CSAF document: {path}')
    doc = obj['document']
    tracking = doc.get('tracking')
    if not isinstance(tracking, dict) or not tracking.get('id') or not tracking.get('current_release_date'):
        raise ValueError(f'Missing CSAF tracking fields: {path}')
    return obj


def atomic_bytes(path, raw):
    path.parent.mkdir(parents=True, exist_ok=True)
    temp = path.with_name(path.name + '.tmp')
    try:
        temp.write_bytes(raw)
        os.replace(temp, path)
    finally:
        temp.unlink(missing_ok=True)


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--since-year', type=int, default=2017)
    p.add_argument('--refresh-old', type=int, default=300, help='Rotate old documents to check changed advisories each run')
    p.add_argument('--full', action='store_true', help='Re-fetch every indexed document')
    p.add_argument('--delay', type=float, default=0.15)
    p.add_argument('--root', default='.')
    args = p.parse_args()
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
    root = Path(args.root).resolve()
    statefile = root / 'state' / 'sync.json'
    state = json.loads(statefile.read_text()) if statefile.exists() else {}
    index_raw = fetch(BASE + 'index.txt')
    entries = parse_index(index_raw, args.since_year)
    if state.get('count', 0) and len(entries) < max(1, int(state['count'] * .70)):
        raise RuntimeError('Index unexpectedly shrunk by over 30%; refusing unsafe sync')
    current = set(entries)
    stored = set(state.get('entries', []))
    pending = sorted(current - stored)
    refresh_n = max(0, args.refresh_old)
    cursor = int(state.get('cursor', 0))
    old = sorted(current & stored)
    revisit = []
    if old and refresh_n:
        revisit = [old[(cursor + n) % len(old)] for n in range(min(refresh_n, len(old)))]
    missing = [item for item in entries if not (root / 'data' / 'csaf' / item).is_file()]
    targets = entries if args.full else sorted(set(pending) | set(missing) | set(revisit))
    LOG.info('Index: %s entries, new=%s missing=%s refreshing=%s total downloads=%s', len(entries), len(pending), len(missing), len(revisit), len(targets))
    updates = 0
    for i, item in enumerate(targets, 1):
        dest = root / 'data' / 'csaf' / item
        raw = fetch(BASE + item)
        validate_document(raw, item)
        if not dest.is_file() or hashlib.sha256(raw).digest() != hashlib.sha256(dest.read_bytes()).digest():
            atomic_bytes(dest, raw)
            updates += 1
        if i % 100 == 0:
            LOG.info('Processed %s/%s (new/updated %s)', i, len(targets), updates)
        if args.delay and i < len(targets):
            time.sleep(args.delay)
    # Apply removals only after ALL downloads succeeded. No state changes on partial failures.
    removed = stored - current
    for item in removed:
        (root / 'data' / 'csaf' / item).unlink(missing_ok=True)
    atomic_bytes(root / 'data' / 'index.txt', index_raw)
    new_state = {
        'source': BASE, 'since_year': args.since_year,
        'last_success_utc': datetime.now(timezone.utc).isoformat(),
        'index_sha256': hashlib.sha256(index_raw).hexdigest(),
        'count': len(entries), 'cursor': (cursor + len(revisit)) % max(len(old), 1),
        'entries': entries,
    }
    atomic_bytes(statefile, (json.dumps(new_state, ensure_ascii=False, sort_keys=True, indent=2) + '\n').encode())
    LOG.info('Done: %s downloaded, %s changed, %s removed', len(targets), updates, len(removed))

if __name__ == '__main__':
    try:
        main()
    except Exception as exc:
        LOG.exception('Sync failed: %s', exc)
        sys.exit(1)