#!/usr/bin/env python3
"""Conservative SQLite index from CVEProject, NVD and OSV raw JSON.
Do not mistake this index for a complete vulnerability matching engine.
"""
import argparse
import hashlib
import json
import os
import re
import shutil
import sqlite3
import subprocess
import tempfile
from datetime import datetime, timezone
from pathlib import Path

CVE = re.compile(r'^CVE-\d{4}-\d{4,}$', re.I)
FILES = ('vulnerabilities','aliases','affected_packages','cpe_matches')


def canonical(obj):
    return json.dumps(obj, ensure_ascii=False, sort_keys=True, separators=(',', ':'))


def sha(s):
    return hashlib.sha256(s.encode('utf-8')).hexdigest()[:32]


def date(value):
    return value if isinstance(value, str) and value else None


def get_git_revision(path):
    return subprocess.check_output(['git', '-C', str(path), 'rev-parse', 'HEAD'], text=True).strip()


def input_files(repo):
    d = repo / 'data'
    if not d.exists():
        return

    # Nur tatsächliche CVE-Datensätze einlesen
    if repo.name == 'vuln-cve':
        yield from sorted(d.rglob('CVE-*.json'))
        return

    # Andere Quellen unverändert verarbeiten
    yield from sorted(d.rglob('*.json'))


def upsert_vuln(db, vid, published=None, modified=None, state=None, summary=None):
    if not vid or not (CVE.match(vid) or vid.startswith(('GHSA-','OSV-','PYSEC-','RUSTSEC-','GO-','DSA-','USN-'))):
        return
    db.execute('''INSERT INTO vulnerabilities(id,published,modified,state,summary) VALUES(?,?,?,?,?)
      ON CONFLICT(id) DO UPDATE SET
      published=COALESCE(excluded.published,vulnerabilities.published),
      modified=COALESCE(excluded.modified,vulnerabilities.modified),
      state=COALESCE(excluded.state,vulnerabilities.state),
      summary=COALESCE(excluded.summary,vulnerabilities.summary)''',
      (vid, published, modified, state, summary[:500] if isinstance(summary,str) else None))


def process_cve(db, obj):
    if not isinstance(obj, dict):
      raise ValueError(
        f"Unexpected CVE JSON type: {type(obj).__name__}; "
        f"sample={str(obj)[:300]}"
    )
    m = obj.get('cveMetadata')
    if not isinstance(m,dict):
        return False
    vid=m.get('cveId')
    if not isinstance(vid,str) or not CVE.match(vid):
        return False
    upsert_vuln(db,vid,m.get('datePublished'),m.get('dateUpdated'),m.get('state'))
    return True


def walk_cpe(node, allow_simple=True):
    """Only emit independent leaf matches for simple non-negated OR nodes.
    Complex AND/NOT expressions must remain unsupported rather than cause false positives.
    """
    if not isinstance(node,dict) or node.get('negate'):
        return
    if node.get('operator','OR').upper() != 'OR':
        return
    matches=node.get('cpeMatch') or []
    if isinstance(matches,list):
        for match in matches:
            if isinstance(match,dict):
                yield match
    for child in node.get('children') or []:
        yield from walk_cpe(child)


def process_nvd(db,obj):
    cve=obj.get('cve')
    if not isinstance(cve,dict):
        return False
    vid=cve.get('id')
    if not isinstance(vid,str) or not CVE.match(vid):
        return False
    upsert_vuln(db,vid,cve.get('published'),cve.get('lastModified'),cve.get('vulnStatus'))
    db.execute('DELETE FROM cpe_matches WHERE vulnerability_id=? AND source=?',(vid,'vuln-nvd'))
    for node in cve.get('configurations') or []:
        # NVD root configurations may contain nodes; a configuration with several
        # nodes can be an AND and is not safe to flatten.
        if not isinstance(node,dict) or node.get('operator','OR').upper() != 'OR' or node.get('negate'):
            continue
        for n in node.get('nodes') or []:
            for match in walk_cpe(n):
                criteria=match.get('criteria')
                if not isinstance(criteria,str):
                    continue
                vals=(criteria, int(bool(match.get('vulnerable',False))),
                      match.get('versionStartIncluding'), match.get('versionStartExcluding'),
                      match.get('versionEndIncluding'),match.get('versionEndExcluding'))
                key=sha(canonical((vid,vals)))
                db.execute('''INSERT OR REPLACE INTO cpe_matches
                  (id,vulnerability_id,criteria,vulnerable,version_start_including,version_start_excluding,version_end_including,version_end_excluding,source)
                  VALUES(?,?,?,?,?,?,?,?,?)''',(key,vid,*vals,'vuln-nvd'))
    return True


def process_osv(db,obj):
    vid=obj.get('id')
    if not isinstance(vid,str) or not vid:
        return False
    # Only capture OSV ids with a known prefix or valid CVE.
    if not (CVE.match(vid) or vid.startswith(('GHSA-','OSV-','PYSEC-','RUSTSEC-','GO-','DSA-','USN-'))):
        return False
    upsert_vuln(db,vid,date(obj.get('published')),date(obj.get('modified')),
                'WITHDRAWN' if obj.get('withdrawn') else None,obj.get('summary'))
    db.execute('DELETE FROM affected_packages WHERE vulnerability_id=? AND source=?',(vid,'vuln-osv'))
    db.execute('DELETE FROM aliases WHERE vulnerability_id=?',(vid,))
    for a in obj.get('aliases') or []:
        if isinstance(a,str) and a:
            db.execute('INSERT OR IGNORE INTO aliases(alias,vulnerability_id) VALUES(?,?)',(a,vid))
    for i,a in enumerate(obj.get('affected') or []):
        if not isinstance(a,dict):
            continue
        pkg=a.get('package') or {}
        if not isinstance(pkg,dict):
            continue
        eco=pkg.get('ecosystem'); name=pkg.get('name')
        ranges=a.get('ranges') or []
        versions=a.get('versions') or []
        if not isinstance(ranges,list) or not isinstance(versions,list):
            continue
        if not isinstance(eco,str) or not isinstance(name,str):
            continue
        # Preserve version/range event semantics; no lexicographic comparisons.
        key=sha(canonical((vid,i,eco,name,ranges,versions)))
        db.execute('INSERT OR REPLACE INTO affected_packages VALUES(?,?,?,?,?,?,?)',
                   (key,vid,eco,name,canonical(ranges),canonical(versions),'vuln-osv'))
    return True


def stream_source(db,name,path):
    count=0; bad=0
    for file in input_files(path):
        try:
            with file.open('r',encoding='utf-8') as f:
                obj=json.load(f)
            if name=='vuln-cve':
                accepted=process_cve(db,obj)
            elif name=='vuln-nvd':
                accepted=process_nvd(db,obj)
            elif name=='vuln-osv':
                accepted=process_osv(db,obj)
            else:
                accepted=False
            count+=int(accepted)
        except (ValueError,UnicodeError,TypeError,sqlite3.Error) as exc:
            bad+=1
            if bad<=5:
                print(f'INVALID {name} {file}: {exc}',flush=True)
    if bad:
        raise RuntimeError(f'{name}: {bad} invalid records; database not published')
    return count


def dump_table(db,table):
    # memory efficient cursor; a canonical row object for deltas
    cols=[r[1] for r in db.execute(f'PRAGMA table_info({table})')]
    for row in db.execute(f'SELECT * FROM {table} ORDER BY 1'):
        yield row[0],dict(zip(cols,row))


def make_delta(old_path,new_db,path):
    old=sqlite3.connect(old_path)
    with path.open('w',encoding='utf-8') as f:
        for table in FILES:
            prev=dump_table(old,table); curr=dump_table(new_db,table)
            a=next(prev,None); b=next(curr,None)
            while a is not None or b is not None:
                if a is not None and (b is None or a[0]<b[0]):
                    f.write(canonical({'table':table,'op':'delete','key':a[0]})+'\n'); a=next(prev,None)
                elif b is not None and (a is None or b[0]<a[0]):
                    f.write(canonical({'table':table,'op':'upsert','row':b[1]})+'\n'); b=next(curr,None)
                else:
                    if a[1]!=b[1]:
                        f.write(canonical({'table':table,'op':'upsert','row':b[1]})+'\n')
                    a=next(prev,None);b=next(curr,None)
    old.close()


def compress(src,out):
    subprocess.run(['zstd','-q','-f','-T1','-19',str(src),'-o',str(out)],check=True)


def main():
    ap=argparse.ArgumentParser()
    ap.add_argument('--sources',type=Path,required=True)
    ap.add_argument('--output',type=Path,required=True)
    ap.add_argument('--schema',type=Path,default=Path(__file__).resolve().parents[1]/'sql/schema.sql')
    args=ap.parse_args()
    args.output.mkdir(parents=True,exist_ok=True)
    active=[(p.name,p) for p in sorted(args.sources.iterdir()) if p.is_dir() and (p/'.git').exists()]
    if not active:
        raise SystemExit('No source repositories checked out; nothing published')
    # Fail closed if expected core sources are missing.
    for name in ('vuln-cve','vuln-nvd','vuln-osv'):
        if name not in [n for n,_ in active]:
            raise SystemExit(f'Missing core source: {name}')
    source_revs={name:get_git_revision(path) for name,path in active}
    old_db=args.output/'master.db'
    with tempfile.TemporaryDirectory(prefix='vuln-build-',dir=args.output) as td:
        tmp=Path(td); new=tmp/'master.db'; db=sqlite3.connect(new)
        db.executescript(args.schema.read_text(encoding='utf-8'))
        stats={}
        for name,path in active:
            if name in ('vuln-cve','vuln-nvd','vuln-osv'):
                stats[name]=stream_source(db,name,path)
            else:
                stats[name]=0
            db.execute('INSERT INTO sources VALUES(?,?,?,?)',(name,source_revs[name],int(name in ('vuln-cve','vuln-nvd','vuln-osv')),stats[name]))
            print(f'{name}: indexed={stats[name]} revision={source_revs[name][:12]}',flush=True)
        if not any(stats.get(n) for n in ('vuln-cve','vuln-nvd','vuln-osv')):
            raise RuntimeError('Core source indexes are empty; refusing to publish')
        old_version=None
        if old_db.is_file():
            with sqlite3.connect(old_db) as o:
                row=o.execute("SELECT value FROM meta WHERE key='version'").fetchone()
                old_version=row[0] if row else None
        version=datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%S%fZ')
        for k,v in (('version',version),('schema_version','1')):
            db.execute('INSERT INTO meta VALUES(?,?)',(k,v))
        db.commit(); db.execute('PRAGMA optimize'); db.execute('VACUUM'); db.commit()
        assert db.execute('PRAGMA integrity_check').fetchone()[0]=='ok'
        db.close()
        manifest={'version':version,'schema_version':1,'previous':old_version,'source_commits':source_revs,
                  'indexed_documents':stats,'format':'sqlite+zstd','delta_format':'jsonl+zstd',
                  'note':'Only CVE/NVD/OSV indexed in v1; no complete version-matching engine'}
        compress(new,tmp/'master.db.zst')
        if old_version:
            with sqlite3.connect(new) as nd:
                make_delta(old_db,nd,tmp/'delta.jsonl')
            compress(tmp/'delta.jsonl',tmp/f'delta-{old_version}-to-{version}.jsonl.zst')
        # Atomically replace individual artifacts; manifest written last as commit marker.
        os.replace(new,old_db)
        os.replace(tmp/'master.db.zst',args.output/'master.db.zst')
        if old_version:
            os.replace(tmp/f'delta-{old_version}-to-{version}.jsonl.zst',args.output/f'delta-{old_version}-to-{version}.jsonl.zst')
        with (args.output/'master.db.zst').open('rb') as hf:
            h=hashlib.sha256()
            for chunk in iter(lambda: hf.read(1024*1024),b''):
                h.update(chunk)
            manifest['master_sha256']=h.hexdigest()
        (tmp/'manifest.json').write_text(json.dumps(manifest,indent=2,ensure_ascii=False)+'\n',encoding='utf-8')
        os.replace(tmp/'manifest.json',args.output/'manifest.json')
        print(f'Build complete: {version}',flush=True)

if __name__=='__main__':
    main()