#!/usr/bin/env python3
"""Shared hammer2 image-walk/forge helpers (cribbed from DF-2616/DF-2627 prior art)."""
import struct

MAGIC = 0x48414D3205172011
RADIX_MASK = 0x3F
BREF = 128
INODE_DATA_BLOCKSET = 0x200
SET_COUNT = 4
VOLHDR_STRIDE = 0x40000
T_EMPTY, T_INODE, T_INDIRECT, T_DATA, T_DIRENT = 0, 1, 2, 3, 4
PFSROOT = 0x01

def _mk():
    poly = 0x82F63B78
    t = []
    for n in range(256):
        c = n
        for _ in range(8):
            c = (c >> 1) ^ poly if (c & 1) else (c >> 1)
        t.append(c)
    return t
_T = _mk()

def iscsi_crc32(data):
    crc = 0xFFFFFFFF
    for b in data:
        crc = _T[(crc ^ b) & 0xFF] ^ (crc >> 8)
    return crc ^ 0xFFFFFFFF

assert iscsi_crc32(b"123456789") == 0xE3069283

def bref_parse(buf, off):
    t, methods, copyid, keybits, vradix, flags = struct.unpack_from('<6B', buf, off)
    key, mtid, modtid, doff, utid = struct.unpack_from('<QQQQQ', buf, off + 8)
    namlen = struct.unpack_from('<H', buf, off + 0x38)[0] if t == T_DIRENT else None
    return dict(type=t, methods=methods, flags=flags, key=key, mirror_tid=mtid,
                modify_tid=modtid, data_off=doff, off=off, namlen=namlen)

def recompute_volhdr_crcs(buf):
    c1 = iscsi_crc32(bytes(buf[512:1024]))
    struct.pack_into('<I', buf, 0x1E0 + 6 * 4, c1)
    c0 = iscsi_crc32(bytes(buf[0:508]))
    struct.pack_into('<I', buf, 0x1E0 + 7 * 4, c0)
    cv = iscsi_crc32(bytes(buf[0:0xFFFC]))
    struct.pack_into('<I', buf, 0xFFFC, cv)

def find_volhdrs(img):
    return [off for off in range(0, 4 * VOLHDR_STRIDE, VOLHDR_STRIDE)
            if struct.unpack_from('<Q', img, off)[0] == MAGIC]

def inode_name(img, iblk):
    nlen = struct.unpack_from('<H', img, iblk + 0x80)[0]
    nlen = min(max(nlen, 0), 255)
    return bytes(img[iblk + 0x100:iblk + 0x100 + max(nlen, 1)]).split(b'\0')[0]

def find_sroot(img):
    """First volhdr + the sroot inode bref + block."""
    vols = find_volhdrs(img)
    assert vols, "no volume header"
    v0 = vols[0]
    for i in range(SET_COUNT):
        br = bref_parse(img, v0 + INODE_DATA_BLOCKSET + i * BREF)
        if br['data_off']:
            return vols, br
    raise AssertionError("no sroot bref in volhdr")

def sroot_children(img, sbr):
    """INODE/other brefs in the sroot inode block's blockset (direct slots)."""
    sblk = sbr['data_off'] & ~RADIX_MASK
    out = []
    for i in range(SET_COUNT):
        br = bref_parse(img, sblk + INODE_DATA_BLOCKSET + i * BREF)
        if br['type'] != T_EMPTY:
            out.append(br)
    return sblk, out

def find_pfs_bref(img, sbr, want_name):
    """Locate the PFS inode bref (direct or under indirects) by filename."""
    sblk, _ = sroot_children(img, sbr)
    for i in range(SET_COUNT):
        br = bref_parse(img, sblk + INODE_DATA_BLOCKSET + i * BREF)
        if br['type'] == T_INODE and (br['data_off'] & RADIX_MASK):
            if inode_name(img, br['data_off'] & ~RADIX_MASK) == want_name:
                return br
        elif br['type'] == T_INDIRECT and (br['data_off'] & RADIX_MASK):
            r = _walk_ind(img, br['data_off'] & ~RADIX_MASK,
                          br['data_off'] & RADIX_MASK, want_name, 0)
            if r:
                return r
    return None

def _walk_ind(img, blk, radix, want_name, depth):
    if depth > 4:
        return None
    nslots = min((1 << radix) // BREF, 512)
    for j in range(nslots):
        br = bref_parse(img, blk + j * BREF)
        if br['type'] == T_EMPTY:
            continue
        if br['type'] == T_INODE and (br['data_off'] & RADIX_MASK):
            if inode_name(img, br['data_off'] & ~RADIX_MASK) == want_name:
                return br
        elif br['type'] == T_INDIRECT and (br['data_off'] & RADIX_MASK):
            r = _walk_ind(img, br['data_off'] & ~RADIX_MASK,
                          br['data_off'] & RADIX_MASK, want_name, depth + 1)
            if r:
                return r
    return None