#!/usr/bin/env python3
"""Wrap the FAT32 boot volume in a hybrid ISO — the flashable danos release image.

Mirrors tools/make-fat-image.py in spirit: pure Python 3 standard library, no
external tools (no xorriso / mkisofs / isohybrid). The output is one file that
boots both ways release media is consumed:

  * Flashed raw to a USB stick (Etcher, dd): the ISO's system area carries an
    MBR whose single partition (type 0xEF, "EFI System") points at the FAT32
    image embedded in the ISO, so UEFI firmware finds the ESP and runs
    \\EFI\\BOOT\\BOOTX64.efi off it.
  * Burned to optical media: an El Torito boot catalog with an EFI platform
    entry points at the same embedded FAT image.

The ISO9660 filesystem itself is minimal but valid — a primary volume
descriptor, the El Torito boot record, path tables, and a root directory that
lists the boot catalog and the FAT image — so inspection tools can open it.

    make-iso-image.py <out.iso> <esp.img>
    make-iso-image.py --verify <out.iso>

All timestamp fields are zero ("not specified") so the build is reproducible.
"""

import struct
import sys

ISO_SECTOR = 2048

# Fixed layout, in ISO sectors (LBA). Sectors 0-15 are the system area (the
# hybrid MBR lives in its first 512 bytes); volume descriptors start at 16.
PVD_LBA = 16          # primary volume descriptor
BOOT_RECORD_LBA = 17  # El Torito boot record volume descriptor
TERMINATOR_LBA = 18   # volume descriptor set terminator
PATH_TABLE_L_LBA = 19
PATH_TABLE_M_LBA = 20
ROOT_DIR_LBA = 21     # root directory (one sector holds our four records)
CATALOG_LBA = 22      # El Torito boot catalog
ESP_LBA = 23          # the embedded FAT32 image starts here

MBR_PARTITION_TYPE_ESP = 0xEF


def both16(value):
    """ISO9660 both-byte-order encoding: little-endian then big-endian."""
    return struct.pack("<H", value) + struct.pack(">H", value)


def both32(value):
    return struct.pack("<I", value) + struct.pack(">I", value)


def directory_record(identifier, lba, size, flags):
    length = 33 + len(identifier)
    if length % 2:
        length += 1  # records are padded to even length
    record = bytearray(length)
    record[0] = length
    record[2:10] = both32(lba)
    record[10:18] = both32(size)
    # record[18:25] is the recording date; zero = unspecified (reproducible).
    record[25] = flags  # 0x02 = directory
    record[28:32] = both16(1)  # volume sequence number
    record[32] = len(identifier)
    record[33:33 + len(identifier)] = identifier
    return bytes(record)


def primary_volume_descriptor(total_sectors, path_table_size):
    sector = bytearray(ISO_SECTOR)
    sector[0] = 1  # type: primary
    sector[1:6] = b"CD001"
    sector[6] = 1  # version
    sector[8:40] = b"DANOS".ljust(32)  # system identifier
    sector[40:72] = b"DANOS".ljust(32)  # volume identifier
    sector[80:88] = both32(total_sectors)
    sector[120:124] = both16(1)  # volume set size
    sector[124:128] = both16(1)  # volume sequence number
    sector[128:132] = both16(ISO_SECTOR)  # logical block size
    sector[132:140] = both32(path_table_size)
    sector[140:144] = struct.pack("<I", PATH_TABLE_L_LBA)
    sector[148:152] = struct.pack(">I", PATH_TABLE_M_LBA)
    sector[156:190] = directory_record(b"\x00", ROOT_DIR_LBA, ISO_SECTOR, 0x02)
    sector[190:318] = b" " * 128  # volume set identifier
    sector[318:446] = b" " * 128  # publisher
    sector[446:574] = b" " * 128  # data preparer
    sector[574:702] = b"DANOS MAKE-ISO-IMAGE".ljust(128)  # application
    sector[702:739] = b" " * 37  # copyright file
    sector[739:776] = b" " * 37  # abstract file
    sector[776:813] = b" " * 37  # bibliographic file
    unspecified_date = b"0" * 16 + b"\x00"
    for offset in (813, 830, 847, 864):  # creation/modification/expiry/effective
        sector[offset:offset + 17] = unspecified_date
    sector[881] = 1  # file structure version
    return bytes(sector)


def boot_record_descriptor():
    sector = bytearray(ISO_SECTOR)
    sector[0] = 0  # type: boot record
    sector[1:6] = b"CD001"
    sector[6] = 1
    sector[7:39] = b"EL TORITO SPECIFICATION".ljust(32, b"\x00")
    sector[71:75] = struct.pack("<I", CATALOG_LBA)
    return bytes(sector)


def terminator_descriptor():
    sector = bytearray(ISO_SECTOR)
    sector[0] = 255
    sector[1:6] = b"CD001"
    sector[6] = 1
    return bytes(sector)


def path_table(byte_order):
    # A single entry: the root directory.
    return (struct.pack("BB", 1, 0)
            + struct.pack(byte_order + "I", ROOT_DIR_LBA)
            + struct.pack(byte_order + "H", 1)
            + b"\x00\x00")  # identifier 0x00 + pad to even


def root_directory(esp_size):
    # Records must be sorted by identifier; BOOT.CAT < EFI.IMG holds.
    entries = (directory_record(b"\x00", ROOT_DIR_LBA, ISO_SECTOR, 0x02)
               + directory_record(b"\x01", ROOT_DIR_LBA, ISO_SECTOR, 0x02)
               + directory_record(b"BOOT.CAT;1", CATALOG_LBA, ISO_SECTOR, 0)
               + directory_record(b"EFI.IMG;1", ESP_LBA, esp_size, 0))
    return entries + b"\x00" * (ISO_SECTOR - len(entries))


def boot_catalog(esp_size):
    # Validation entry: EFI platform (0xEF), checksummed so its 16-bit words sum
    # to zero, closed by the 0x55AA key bytes.
    validation = bytearray(32)
    validation[0] = 0x01
    validation[1] = 0xEF
    validation[4:28] = b"danos".ljust(24, b"\x00")
    validation[30] = 0x55
    validation[31] = 0xAA
    checksum = (-sum(struct.unpack("<16H", validation))) & 0xFFFF
    validation[28:30] = struct.pack("<H", checksum)
    # Initial/default entry: bootable, no emulation, image at ESP_LBA. The
    # sector-count field is 16-bit (units of 512 bytes) so it can't span a large
    # ESP; UEFI firmware sizes the FAT filesystem from its own BPB, and the
    # image's boot files sit well inside the capped span regardless.
    default = bytearray(32)
    default[0] = 0x88  # bootable
    default[1] = 0x00  # no emulation
    sector_count = min(0xFFFF, esp_size // 512)
    default[6:8] = struct.pack("<H", sector_count)
    default[8:12] = struct.pack("<I", ESP_LBA)
    catalog = bytes(validation) + bytes(default)
    return catalog + b"\x00" * (ISO_SECTOR - len(catalog))


def hybrid_mbr(esp_size):
    """The system-area MBR that makes the ISO flashable: one ESP partition."""
    mbr = bytearray(512)
    mbr[440:444] = b"dano"  # disk signature (fixed: reproducible builds)
    start_lba = ESP_LBA * (ISO_SECTOR // 512)
    partition = struct.pack(
        "<B3sB3sII",
        0x80,                      # status: active (harmless; helps picky firmware)
        b"\xFE\xFF\xFF",           # CHS start: maxed out, LBA is authoritative
        MBR_PARTITION_TYPE_ESP,    # type: EFI System
        b"\xFE\xFF\xFF",           # CHS end
        start_lba,
        esp_size // 512,
    )
    mbr[446:462] = partition
    mbr[510] = 0x55
    mbr[511] = 0xAA
    return bytes(mbr)


def build(out_path, esp_path):
    with open(esp_path, "rb") as handle:
        esp = handle.read()
    if len(esp) % ISO_SECTOR:
        esp += b"\x00" * (ISO_SECTOR - len(esp) % ISO_SECTOR)
    esp_sectors = len(esp) // ISO_SECTOR
    total_sectors = ESP_LBA + esp_sectors

    table_l = path_table("<")
    image = bytearray(total_sectors * ISO_SECTOR)
    image[0:512] = hybrid_mbr(len(esp))
    image[PVD_LBA * ISO_SECTOR:(PVD_LBA + 1) * ISO_SECTOR] = \
        primary_volume_descriptor(total_sectors, len(table_l))
    image[BOOT_RECORD_LBA * ISO_SECTOR:(BOOT_RECORD_LBA + 1) * ISO_SECTOR] = \
        boot_record_descriptor()
    image[TERMINATOR_LBA * ISO_SECTOR:(TERMINATOR_LBA + 1) * ISO_SECTOR] = \
        terminator_descriptor()
    image[PATH_TABLE_L_LBA * ISO_SECTOR:PATH_TABLE_L_LBA * ISO_SECTOR + len(table_l)] = table_l
    table_m = path_table(">")
    image[PATH_TABLE_M_LBA * ISO_SECTOR:PATH_TABLE_M_LBA * ISO_SECTOR + len(table_m)] = table_m
    image[ROOT_DIR_LBA * ISO_SECTOR:(ROOT_DIR_LBA + 1) * ISO_SECTOR] = root_directory(len(esp))
    image[CATALOG_LBA * ISO_SECTOR:(CATALOG_LBA + 1) * ISO_SECTOR] = boot_catalog(len(esp))
    image[ESP_LBA * ISO_SECTOR:] = esp

    with open(out_path, "wb") as handle:
        handle.write(image)
    print(f"make-iso-image: wrote {out_path} "
          f"({total_sectors * ISO_SECTOR // (1024 * 1024)} MiB hybrid ISO, "
          f"ESP at LBA {ESP_LBA}, {esp_sectors} sectors)")


def verify(path):
    with open(path, "rb") as handle:
        data = handle.read()
    # The hybrid MBR (the Etcher/dd boot path).
    if data[510] != 0x55 or data[511] != 0xAA:
        sys.exit("verify: missing MBR 0x55AA signature")
    status, _, part_type, _, part_start, part_sectors = \
        struct.unpack_from("<B3sB3sII", data, 446)
    if part_type != MBR_PARTITION_TYPE_ESP:
        sys.exit(f"verify: MBR partition type 0x{part_type:02X}, expected 0xEF (ESP)")
    # The ISO9660 descriptors (the optical boot path).
    if data[PVD_LBA * ISO_SECTOR + 1:PVD_LBA * ISO_SECTOR + 6] != b"CD001":
        sys.exit("verify: no primary volume descriptor")
    boot_record = data[BOOT_RECORD_LBA * ISO_SECTOR:(BOOT_RECORD_LBA + 1) * ISO_SECTOR]
    if not boot_record.startswith(b"\x00CD001") or \
            not boot_record[7:30].startswith(b"EL TORITO SPECIFICATION"):
        sys.exit("verify: no El Torito boot record")
    catalog_lba = struct.unpack_from("<I", boot_record, 71)[0]
    catalog = data[catalog_lba * ISO_SECTOR:(catalog_lba + 1) * ISO_SECTOR]
    if catalog[0] != 0x01 or catalog[1] != 0xEF or catalog[30:32] != b"\x55\xAA":
        sys.exit("verify: boot catalog validation entry is not an EFI entry")
    if sum(struct.unpack("<16H", catalog[0:32])) & 0xFFFF != 0:
        sys.exit("verify: boot catalog validation checksum is wrong")
    if catalog[32] != 0x88:
        sys.exit("verify: default catalog entry is not bootable")
    boot_lba = struct.unpack_from("<I", catalog, 40)[0]
    # Both paths must agree on where the ESP lives, and it must be a FAT32 image.
    if boot_lba * (ISO_SECTOR // 512) != part_start:
        sys.exit(f"verify: catalog boot image (LBA {boot_lba}) and MBR partition "
                 f"(sector {part_start}) disagree")
    esp = data[boot_lba * ISO_SECTOR:]
    if len(esp) < part_sectors * 512:
        sys.exit("verify: MBR partition extends past the end of the file")
    if esp[510] != 0x55 or esp[511] != 0xAA or esp[82:90] != b"FAT32   ":
        sys.exit("verify: embedded image is not a FAT32 boot volume")
    print(f"verify: {path} is a hybrid ISO — MBR ESP partition (sector {part_start}, "
          f"{part_sectors} sectors, active={status == 0x80}) and El Torito EFI entry "
          f"both point at the embedded FAT32 image")


def main(argv):
    if len(argv) == 3 and argv[1] == "--verify":
        verify(argv[2])
        return 0
    if len(argv) != 3:
        sys.exit("usage: make-iso-image.py <out.iso> <esp.img>\n"
                 "       make-iso-image.py --verify <out.iso>")
    build(argv[1], argv[2])
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))