#!/usr/bin/env python3
"""
scan_bdb.py -- find Berkeley DB pages written into files that must not contain them
(blocks/blk*.dat, blocks/rev*.dat).  Stand-in for the production incident tool.

BDB writes page P of a database at file offset P*pagesize through pwrite(fd, ...).  If
fd has been closed and its number reused by a block file, the page lands in the block
file at exactly P*pagesize.  So at every pagesize-aligned offset we decode the PAGE
header (db_page.h, all little-endian on x86):
    0  u32 lsn.file     4  u32 lsn.offset   8  u32 pgno   12 u32 prev_pgno
   16  u32 next_pgno   20  u16 entries     22  u16 hf_offset
   24  u8  level       25  u8  type
and a metadata page (pgno 0 / any meta type) also carries magic at 12 and pagesize at 20.

A hit needs ALL of:
  * pgno == offset // pagesize                       (the self-locating property)
  * type is a real BDB page type (1..16 except 0)
  * level consistent with type (overflow/meta/hash: 0; btree leaf: 1; internal: 2..32)
  * hf_offset <= pagesize; entries <= pagesize/2
  * for overflow (7): entries (refcount) >= 1 and 0 < hf_offset (data length) <= pagesize-26
Strength: "STRONG" if lsn == (0,1) (pages rewritten by lsn_reset -- the production
signature) or a meta page with a valid magic, else "weak".
Also reports any aligned BDB meta magic (btree 0x053162, hash 0x061561, queue 0x042253,
heap 0x074582) even when pgno does not match.

usage: scan_bdb.py [--pagesize 4096] [--all-offsets] files...
exit 0 = clean, 2 = hits.
"""
import argparse, struct, sys

TYPES = {1: "duplicate", 2: "hash_unsorted", 3: "ibtree", 4: "irecno", 5: "lbtree", 6: "lrecno",
         7: "overflow", 8: "hashmeta", 9: "btreemeta", 10: "qammeta", 11: "qamdata", 12: "ldup",
         13: "hash", 14: "heapmeta", 15: "heap", 16: "iheap"}
MAGIC = {0x053162: "btree", 0x061561: "hash", 0x042253: "queue", 0x074582: "heap"}
META = {8, 9, 10, 14}


def level_ok(t, lvl):
    if t in (5, 6, 12):           # leaf pages
        return lvl == 1
    if t in (3, 4):               # internal pages
        return 2 <= lvl <= 32
    return lvl == 0               # overflow / meta / hash / queue / heap


def check(buf, off, ps):
    if len(buf) < 26:
        return None
    lf, lo, pgno, prv, nxt, ent, hfo, lvl, typ = struct.unpack_from("<IIIIIHHBB", buf, 0)
    if typ not in TYPES or pgno != off // ps or not level_ok(typ, lvl):
        return None
    if typ in META:
        magic, ver, psz = struct.unpack_from("<III", buf, 12)
        if magic not in MAGIC:
            return None
        return ("STRONG", f"{TYPES[typ]} magic={MAGIC[magic]} ver={ver} pagesize={psz} lsn=({lf},{lo})")
    if hfo > ps or ent > ps // 2:
        return None
    if typ == 7 and not (ent >= 1 and 0 < hfo <= ps - 26):
        return None
    strength = "STRONG" if (lf, lo) == (0, 1) else "weak"
    return (strength, f"{TYPES[typ]} lsn=({lf},{lo}) prev={prv} next={nxt} entries={ent} hf_offset={hfo} level={lvl}")


def scan(path, ps):
    hits, magic_only = [], []
    with open(path, "rb") as f:
        off = 0
        while True:
            buf = f.read(ps)
            if not buf:
                break
            r = check(buf, off, ps)
            if r:
                hits.append((off, off // ps) + r)
            elif len(buf) >= 16:
                m = struct.unpack_from("<I", buf, 12)[0]
                if m in MAGIC and buf[25] in META:
                    magic_only.append((off, MAGIC[m]))
            off += len(buf)
    return hits, magic_only, off


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pagesize", type=int, default=4096)
    ap.add_argument("--quiet", action="store_true")
    ap.add_argument("files", nargs="+")
    a = ap.parse_args()
    total = strong = 0
    for p in a.files:
        try:
            hits, mo, size = scan(p, a.pagesize)
        except OSError as e:
            print(f"{p}: ERROR {e}")
            continue
        s = sum(1 for h in hits if h[2] == "STRONG")
        total += len(hits); strong += s
        if hits or mo:
            print(f"{p}: {len(hits)} BDB page(s) ({s} strong), {len(mo)} stray meta magic, size={size}")
            if not a.quiet:
                for h in hits[:20]:
                    print(f"    off={h[0]} pgno={h[1]} {h[2]} {h[3]}")
                for m in mo[:5]:
                    print(f"    meta-magic off={m[0]} {m[1]}")
        else:
            print(f"{p}: clean ({size} bytes)")
    print(f"TOTAL BDB pages found: {total} ({strong} strong)")
    return 2 if total else 0


if __name__ == "__main__":
    sys.exit(main())