#!/usr/bin/env python3
"""
DF-2642 offline media walker.

Walks a hammer2 image (raw copy taken while mounted/at the wall) from the
volume header -> sroot inode -> PFS 'DATA' root inode (through INDIRECT
blocks as needed) -> victim file inodes -> their blockset DATA brefs.

For every victim (identified by the content of its key-0 DATA block,
'DF2642-Vnn-B0-OLD-'), reports whether the key-65536 (block1) DATA bref
is STILL PRESENT in the on-media topology and what the media content at
its data_off is:

    Vnn: B1-BREF PRESENT, media='DF2642-Vnn-B1-OLD-...'   -> STALE (bug)
    Vnn: B1-bref absent/EMPTY                              -> deleted (ok)

Usage: walk2642.py <image>
"""
import struct
import sys

RADIX_MASK = 0x3F
BREF = 128
SET_COUNT = 4
INODE_DATA_BLOCKSET = 0x200

BREF_EMPTY, BREF_INODE, BREF_INDIRECT, BREF_DATA, BREF_DIRENT = 0, 1, 5, 3, 4

VOLHDR_COPIES = [0x0, 0x40000, 0x80000, 0xC0000]
VOLHDR_MAGIC = 0x48414D3205172011


def bref_parse(buf, off):
    t, methods = struct.unpack_from('<2B', buf, off)
    key, mirror_tid = struct.unpack_from('<QQ', buf, off + 8)
    data_off = struct.unpack_from('<Q', buf, off + 0x20)[0]
    return dict(type=t, methods=methods, key=key, mirror_tid=mirror_tid,
                data_off=data_off, off=off)


def block(img, data_off):
    base = data_off & ~RADIX_MASK
    radix = data_off & RADIX_MASK
    if radix == 0 or base == 0:
        return None
    return img[base:base + (1 << radix)]


def inode_meta(blk):
    # hammer2_inode_data_t: op_flags etc.; meta.inum @ +0x58 per DF-2640 notes
    inum = struct.unpack_from('<Q', blk, 0x58)[0]
    op_flags = blk[0x51]
    ftype = blk[0x50]
    size = struct.unpack_from('<Q', blk, 0x60)[0]
    return inum, op_flags, ftype, size


def collect_brefs(img, inode_blk):
    """yield brefs of an inode's direct blockset (8 slots)"""
    for i in range(SET_COUNT):
        off = INODE_DATA_BLOCKSET + i * BREF
        yield bref_parse(inode_blk, off)


def walk_indirect(img, ind_blk, depth=0):
    """yield brefs stored in an indirect block (recursively)"""
    count = len(ind_blk) // BREF
    for i in range(count):
        br = bref_parse(ind_blk, i * BREF)
        if br['type'] == BREF_EMPTY:
            continue
        yield br
        if br['type'] == BREF_INDIRECT and depth < 8:
            sub = block(img, br['data_off'])
            if sub:
                yield from walk_indirect(img, sub, depth + 1)


def main():
    img = open(sys.argv[1], 'rb').read()
    assert len(img) >= 0x100000, "image too small"

    # pick the volume header copy with the highest mirror_tid (newest)
    best, best_tid = None, -1
    for off in VOLHDR_COPIES:
        if struct.unpack_from('<Q', img, off)[0] != VOLHDR_MAGIC:
            continue
        sroot_bset = off + 0x200
        for i in range(SET_COUNT):
            br = bref_parse(img, sroot_bset + i * BREF)
            if br['type'] == BREF_INODE and br['data_off'] & RADIX_MASK:
                tid = br['mirror_tid']
                if tid > best_tid:
                    best, best_tid = br, tid
    assert best, "no sroot inode bref in any volhdr copy"
    print("[walk] sroot INODE bref data_off=%#x mirror_tid=%#x"
          % (best['data_off'], best_tid))

    sroot_blk = block(img, best['data_off'])
    assert sroot_blk, "sroot block unreadable"

    # find PFS root INODE brefs under the sroot (through indirects)
    pfs_roots = []
    def scan_inode_children(inode_blk, out, depth=0):
        for br in collect_brefs(img, inode_blk):
            if br['type'] == BREF_INODE and br['data_off'] & RADIX_MASK:
                out.append(br)
            elif br['type'] == BREF_INDIRECT:
                sub = block(img, br['data_off'])
                if sub:
                    for br2 in walk_indirect(img, sub):
                        if br2['type'] == BREF_INODE and \
                           br2['data_off'] & RADIX_MASK:
                            out.append(br2)
    scan_inode_children(sroot_blk, pfs_roots)
    print("[walk] %d candidate INODE children under sroot" % len(pfs_roots))

    stale = ok = 0
    file_inodes = []
    for br in pfs_roots:
        pblk = block(img, br['data_off'])
        if pblk is None:
            continue
        scan_inode_children(pblk, file_inodes)
    print("[walk] %d INODE candidates under PFS roots" % len(file_inodes))
    for br in file_inodes:
        blk = block(img, br['data_off'])
        if blk is None:
            continue
        try:
            inum, op_flags, ftype, size = inode_meta(blk)
        except Exception:
            continue
        if ftype != 2 or size != 131072:  # OBJTYPE_REGFILE, 128KB victims
            continue
        # victim: has a key-0 DATA bref with DF2642 pattern
        ident = None
        for d in collect_brefs(img, blk):
            if d['type'] == BREF_DATA and d['key'] == 0:
                dblk = block(img, d['data_off'])
                if dblk and dblk[:6] == b'DF2642':
                    ident = dblk[:18].decode('latin1')
                    break
        if not ident:
            continue
        b1 = None
        for d in collect_brefs(img, blk):
            if d['type'] == BREF_DATA and d['key'] == 65536:
                b1 = d
                break
        if b1 is None:
            print("%s inum=%#x: B1-BREF ABSENT (deleted -> reads as zeros)"
                  % (ident, inum))
            ok += 1
        else:
            dblk = block(img, b1['data_off'])
            head = dblk[:24] if dblk else b'<unreadable>'
            tag = 'STALE' if head[:6] == b'DF2642' else 'WEIRD'
            print("%s inum=%#x: B1-BREF PRESENT data_off=%#x media=%r"
                  " -> %s" % (ident, inum, b1['data_off'], head, tag))
            if tag == 'STALE':
                stale += 1
    print("WALK_DONE stale=%d ok=%d" % (stale, ok))


if __name__ == '__main__':
    main()