provenance.py

"""
provenance.py - Tamper-evident audit ledger for DocSentry.

Every document analysis is logged as a record in a SQLite database. Each
record stores:
    - timestamp
    - document filename + SHA-256
    - verdict band + risk score
    - SHA-256 of the previous ledger entry (hash chain)
    - record_hash: SHA-256 of (timestamp + doc hash + verdict + prev_hash)

The chain construction means that retroactively tampering with any past
record breaks the hash chain - giving DocSentry an audit trail that meets
RBI's tamper-evident record-retention requirements (Master Direction on
KYC, 2016, Para 67).

This is conceptually a baby blockchain: hash-linked records, immutable in
practice, verifiable in O(N).

Public API:
    log_analysis(doc_path, sha256, risk_band, risk_score, extra)
    fetch_ledger(limit=100)
    verify_chain()                    - returns (ok: bool, broken_at: index)
    chain_stats()                     - summary numbers for the UI
    ledger_dataframe()                - pandas DataFrame view
"""

import os
import json
import sqlite3
import hashlib
from datetime import datetime
from pathlib import Path

LEDGER_PATH = Path("provenance.db")


def _conn():
    LEDGER_PATH.parent.mkdir(parents=True, exist_ok=True)
    con = sqlite3.connect(str(LEDGER_PATH))
    con.execute("""CREATE TABLE IF NOT EXISTS ledger (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        ts          TEXT NOT NULL,
        doc_name    TEXT NOT NULL,
        doc_sha256  TEXT NOT NULL,
        risk_band   TEXT,
        risk_score  REAL,
        extra_json  TEXT,
        prev_hash   TEXT,
        record_hash TEXT NOT NULL
    )""")
    con.commit()
    return con


def _hash_record(ts, doc_sha256, risk_band, risk_score, prev_hash):
    payload = f"{ts}|{doc_sha256}|{risk_band}|{risk_score}|{prev_hash}".encode()
    return hashlib.sha256(payload).hexdigest()


def log_analysis(doc_path, sha256, risk_band, risk_score, extra=None):
    """Append a new entry to the provenance ledger and return the new row."""
    ts = datetime.utcnow().isoformat() + "Z"
    con = _conn()
    cur = con.cursor()
    cur.execute("SELECT record_hash FROM ledger ORDER BY id DESC LIMIT 1")
    last = cur.fetchone()
    prev_hash = last[0] if last else "GENESIS"
    rec_hash = _hash_record(ts, sha256, risk_band, risk_score, prev_hash)
    cur.execute("""INSERT INTO ledger (ts, doc_name, doc_sha256, risk_band,
                                       risk_score, extra_json, prev_hash, record_hash)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                (ts, Path(doc_path).name, sha256, risk_band, risk_score,
                 json.dumps(extra or {}, default=str), prev_hash, rec_hash))
    con.commit()
    new_id = cur.lastrowid
    con.close()
    return {"id": new_id, "ts": ts, "record_hash": rec_hash, "prev_hash": prev_hash}


def fetch_ledger(limit=100):
    con = _conn()
    rows = con.execute("""SELECT id, ts, doc_name, doc_sha256, risk_band,
                                 risk_score, prev_hash, record_hash
                          FROM ledger ORDER BY id DESC LIMIT ?""", (limit,)).fetchall()
    con.close()
    return [{
        "id": r[0], "ts": r[1], "doc_name": r[2],
        "doc_sha256": r[3], "risk_band": r[4], "risk_score": r[5],
        "prev_hash": r[6], "record_hash": r[7],
    } for r in rows]


def verify_chain():
    """
    Walk the ledger in order and verify each record_hash matches what
    should have been computed from (ts, doc_sha256, risk_band, risk_score,
    prev_hash). Returns (ok, broken_at).
    """
    con = _conn()
    rows = con.execute("""SELECT id, ts, doc_sha256, risk_band, risk_score,
                                 prev_hash, record_hash FROM ledger ORDER BY id""").fetchall()
    con.close()
    prev = "GENESIS"
    for r in rows:
        rid, ts, doc_hash, band, score, stored_prev, stored_rec = r
        if stored_prev != prev:
            return False, rid
        expected = _hash_record(ts, doc_hash, band, score, prev)
        if expected != stored_rec:
            return False, rid
        prev = stored_rec
    return True, None


def chain_stats():
    con = _conn()
    n = con.execute("SELECT COUNT(*) FROM ledger").fetchone()[0]
    bands = con.execute("SELECT risk_band, COUNT(*) FROM ledger GROUP BY risk_band").fetchall()
    first = con.execute("SELECT ts FROM ledger ORDER BY id LIMIT 1").fetchone()
    last = con.execute("SELECT ts FROM ledger ORDER BY id DESC LIMIT 1").fetchone()
    con.close()
    ok, broken_at = verify_chain()
    return {
        "n_records": n,
        "first_ts": first[0] if first else None,
        "last_ts": last[0] if last else None,
        "by_band": dict(bands),
        "chain_intact": ok,
        "broken_at": broken_at,
    }


def ledger_dataframe(limit=100):
    """Convenience: return ledger as a pandas DataFrame for Streamlit display."""
    import pandas as pd
    rows = fetch_ledger(limit=limit)
    return pd.DataFrame(rows)


def clear_ledger():
    """Wipe the ledger (for demo resets only)."""
    con = _conn()
    con.execute("DELETE FROM ledger")
    con.commit()
    con.close()


if __name__ == "__main__":
    # Smoke test
    clear_ledger()
    log_analysis("doc1.pdf", "a" * 64, "LOW", 0.05)
    log_analysis("doc2.pdf", "b" * 64, "MEDIUM", 0.42)
    log_analysis("doc3.pdf", "c" * 64, "HIGH", 0.74)
    log_analysis("doc4.pdf", "d" * 64, "CRITICAL", 0.91)
    print("Stats:", chain_stats())
    print("\nLast 4 records:")
    for r in fetch_ledger(4):
        print(f"  #{r['id']} {r['ts']} {r['doc_name']:8s} {r['risk_band']:9s} "
              f"hash={r['record_hash'][:12]}... prev={r['prev_hash'][:12]}...")
    ok, broken = verify_chain()
    print(f"\nChain verify: ok={ok}, broken_at={broken}")
    # Simulate tampering
    con = _conn()
    con.execute("UPDATE ledger SET risk_band='LOW' WHERE id=3")
    con.commit()
    con.close()
    ok, broken = verify_chain()
    print(f"After tampering: ok={ok}, broken_at={broken}")
    clear_ledger()
    print("Ledger cleared.")