66 lines
2.6 KiB
Python
66 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Apply delta atomically, rejecting gaps or mismatched base database."""
|
|
import argparse
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import shutil
|
|
import sqlite3
|
|
import tempfile
|
|
import subprocess
|
|
|
|
def fingerprint(conn):
|
|
h = hashlib.sha256()
|
|
for cve_id, row_hash in conn.execute('SELECT cve_id,row_hash FROM cve ORDER BY cve_id'):
|
|
h.update(cve_id.encode() + b'\x00' + row_hash.encode() + b'\n')
|
|
return h.hexdigest()
|
|
|
|
def apply(db, delta):
|
|
db = Path(db)
|
|
with tempfile.TemporaryDirectory(dir=db.parent) as tmp:
|
|
work = Path(tmp) / 'master.db'
|
|
shutil.copy2(db, work)
|
|
conn = sqlite3.connect(work)
|
|
proc = subprocess.Popen(['zstd', '-q', '-dc', str(delta)], stdout=subprocess.PIPE, text=True)
|
|
try:
|
|
lines = proc.stdout
|
|
header = json.loads(next(lines))
|
|
if header.get('format') != 1 or header.get('type') != 'header':
|
|
raise ValueError('Invalid delta header')
|
|
if fingerprint(conn) != header['base']:
|
|
raise ValueError('Delta base mismatch: full master or missing prior delta required')
|
|
conn.execute('BEGIN IMMEDIATE')
|
|
for line in lines:
|
|
op = json.loads(line)
|
|
if op['type'] == 'delete':
|
|
conn.execute('DELETE FROM cve WHERE cve_id=?', (op['id'],))
|
|
elif op['type'] == 'upsert':
|
|
blob = base64.b64decode(op['affected_b64'], validate=True)
|
|
expected = hashlib.sha256(json.dumps([op['id'], op['state'], op['updated']], separators=(',', ':')).encode() + blob).hexdigest()
|
|
if expected != op['hash']:
|
|
raise ValueError('Delta row hash mismatch')
|
|
conn.execute('INSERT OR REPLACE INTO cve VALUES (?,?,?,?,?)', (op['id'], op['state'], op['updated'], blob, op['hash']))
|
|
else:
|
|
raise ValueError('Unknown delta operation')
|
|
if proc.wait() != 0:
|
|
raise ValueError('zstd decompression failed')
|
|
if fingerprint(conn) != header['target']:
|
|
raise ValueError('Delta target mismatch')
|
|
conn.commit()
|
|
finally:
|
|
proc.stdout.close()
|
|
if proc.poll() is None:
|
|
proc.kill()
|
|
proc.wait()
|
|
conn.close()
|
|
os.replace(work, db)
|
|
|
|
if __name__ == '__main__':
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument('--database', required=True, type=Path)
|
|
p.add_argument('--delta', required=True, type=Path)
|
|
a = p.parse_args()
|
|
apply(a.database, a.delta)
|