#!/usr/bin/env python3
"""Store raw CISA KEV JSON and FIRST EPSS daily compressed CSV in Git."""
import csv
import gzip
import io
import json
import os
import re
import tempfile
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SOURCES = {
    'kev': 'https://www.cisa.gov/sites/default/files/feeds/known_exploited_vulnerabilities.json',
    'epss': 'https://epss.empiricalsecurity.com/epss_scores-current.csv.gz',
}

def download(url):
    req = urllib.request.Request(url, headers={'User-Agent': 'vuln-exploits-mirror/1.0'})
    for retry in range(5):
        try:
            with urllib.request.urlopen(req, timeout=180) as resp:
                return resp.read()
        except urllib.error.HTTPError as exc:
            if exc.code not in (408, 429, 500, 502, 503, 504) or retry == 4:
                raise
        except (urllib.error.URLError, TimeoutError):
            if retry == 4:
                raise
        time.sleep(3 * 2**retry)

def atomic_write(path, content):
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.exists() and path.read_bytes() == content:
        return False
    temp = path.with_name(path.name + '.tmp')
    try:
        temp.write_bytes(content)
        os.replace(temp, path)
    finally:
        temp.unlink(missing_ok=True)
    return True

def main():
    kev = download(SOURCES['kev'])
    data = json.loads(kev)
    records = data.get('vulnerabilities')
    if not isinstance(records, list) or len(records) < 1000 or data.get('count') != len(records):
        raise ValueError('Unexpected or truncated CISA KEV JSON')
    if not all(re.fullmatch(r'CVE-\d{4}-\d{4,}', r.get('cveID','')) for r in records):
        raise ValueError('CISA KEV contains invalid CVE IDs')
    epss = download(SOURCES['epss'])
    with gzip.GzipFile(fileobj=io.BytesIO(epss)) as zipped:
        csv_text = io.TextIOWrapper(zipped, encoding='utf-8-sig', newline='')
        reader = csv.reader(line for line in csv_text if not line.startswith('#'))
        header = next(reader)
        if header != ['cve', 'epss', 'percentile']:
            raise ValueError(f'Unexpected EPSS CSV header: {header}')
        rows = sum(1 for row in reader if len(row) == 3 and re.fullmatch(r'CVE-\d{4}-\d{4,}', row[0]))
        if rows < 10000:
            raise ValueError(f'EPSS snapshot too small: {rows}')
    # Only update files after both datasets validated.
    date = datetime.now(timezone.utc).strftime('%Y-%m-%d')
    output = ROOT/'data'
    changed_kev = atomic_write(output/'cisa-kev'/'catalog.json', kev)
    changed_epss = atomic_write(output/'epss'/'current.csv.gz', epss)
    # Daily snapshots are optional: keeping all of them significantly grows Git.
    if os.getenv('EPSS_DAILY_HISTORY', '0') == '1':
        atomic_write(output/'epss'/'history'/f'{date}.csv.gz', epss)
    state = {'version': 1, 'last_success': datetime.now(timezone.utc).isoformat(), 'kev_count': len(records), 'epss_count': rows,
             'epss_history_enabled': os.getenv('EPSS_DAILY_HISTORY', '0') == '1'}
    atomic_write(ROOT/'state'/'sync.json', (json.dumps(state, indent=2, sort_keys=True)+'\n').encode())
    print(f'KEV={len(records)} changed={changed_kev}, EPSS={rows} changed={changed_epss}')

if __name__ == '__main__':
    main()