#!/usr/bin/env python3
# rpi4-efi-bootorder.py - make MontaukOS the Pi UEFI firmware's first boot option
#
# Usage: rpi4-efi-bootorder.py <RPI_EFI.fd> <output.fd> <disk-id> <partition> <start-lba> <sectors>
#        (run by mk-rpi4-sd.sh)
#
# The firmware (github.com/pftf/RPi4, EDK2) keeps its UEFI variables in an
# EDK2 variable store inside RPI_EFI.fd, empty in a release. Left to itself it
# builds the boot order from whatever it enumerates (SD, USB, PXE and HTTP over
# IPv4 and IPv6, ...), appending entries it finds later, so the card can end up
# behind the network entries and booting it means picking it from the boot
# menu. This writes two variables into the empty store:
#
#   Boot0000   "MontaukOS": HD(<partition>, MBR, <disk-id>)/\EFI\BOOT\BOOTAA64.EFI
#   BootOrder  { 0000 }
#
# The device path is the short form that starts at the partition: EDK2 finds
# it by the MBR disk ID and partition number on any block device (connecting
# everything the first time it looks, then caching the full path), so it
# boots the card whether it is in the SD slot or written to a USB drive. The
# firmware appends the entries it creates after it, and never deletes it (it
# only prunes the entries it created itself). The image's MBR must carry the
# same disk ID; montauk-imager.py's `update --firmware` copies it to the card.

import struct
import sys
import uuid

NV_FV_GUID = uuid.UUID("fff12b8d-7696-4c8b-a985-2747075b4f50").bytes_le     # gEfiSystemNvDataFvGuid
AUTH_STORE_GUID = uuid.UUID("aaf32c78-947b-439a-a180-2e144ec37792").bytes_le  # gEfiAuthenticatedVariableGuid
GLOBAL_GUID = uuid.UUID("8be4df61-93ca-11d2-aa0d-00e098032b8c").bytes_le     # gEfiGlobalVariableGuid

STORE_FORMATTED = 0x5A
STORE_HEALTHY = 0xFE
VARIABLE_DATA = 0x55AA
VAR_ADDED = 0x3F
NV_BS_RT = 0x7                  # NON_VOLATILE | BOOTSERVICE_ACCESS | RUNTIME_ACCESS
LOAD_OPTION_ACTIVE = 0x1
LOADER = "\\EFI\\BOOT\\BOOTAA64.EFI"


def fail(msg):
    print(f"rpi4-efi-bootorder: {msg}", file=sys.stderr)
    sys.exit(1)


def find_store(fd):
    """Offsets of the variable store header and of its end."""
    for fv in range(0, len(fd) - 0x48, 0x1000):
        if fd[fv + 0x10:fv + 0x20] == NV_FV_GUID and fd[fv + 0x28:fv + 0x2C] == b"_FVH":
            store = fv + struct.unpack_from("<H", fd, fv + 0x30)[0]
            guid, size, fmt, state = struct.unpack_from("<16sIBB", fd, store)
            if guid != AUTH_STORE_GUID or fmt != STORE_FORMATTED or state != STORE_HEALTHY:
                fail("the NV variable store is not an authenticated EDK2 store")
            return store, store + size
    fail("no NV variable store in the firmware image")


def variable(name, data):
    # AUTHENTICATED_VARIABLE_HEADER (packed, 60 bytes), name, data; the next
    # header starts 4-byte aligned, the gap left erased.
    name = (name + "\0").encode("utf-16-le")
    v = struct.pack("<HBBIQ16sIII16s", VARIABLE_DATA, VAR_ADDED, 0, NV_BS_RT, 0, bytes(16), 0,
                    len(name), len(data), GLOBAL_GUID) + name + data
    return v + b"\xff" * (-len(v) % 4)


def load_option(desc, disk_id, part, start, sectors):
    hd = struct.pack("<BBHIQQ16sBB", 4, 1, 42, part, start, sectors,
                     struct.pack("<I", disk_id).ljust(16, b"\0"), 1, 1)   # MBR, MBR signature
    path = (LOADER + "\0").encode("utf-16-le")
    dp = hd + struct.pack("<BBH", 4, 4, 4 + len(path)) + path + struct.pack("<BBH", 0x7F, 0xFF, 4)
    return struct.pack("<IH", LOAD_OPTION_ACTIVE, len(dp)) + (desc + "\0").encode("utf-16-le") + dp


def main():
    if len(sys.argv) != 7:
        fail("usage: rpi4-efi-bootorder.py <RPI_EFI.fd> <output.fd> <disk-id> <partition> <start-lba> <sectors>")
    src, out = sys.argv[1], sys.argv[2]
    disk_id, part, start, sectors = (int(a, 0) for a in sys.argv[3:])
    if not 0 < disk_id <= 0xFFFFFFFF:
        fail(f"bad MBR disk ID {sys.argv[3]}")

    fd = bytearray(open(src, "rb").read())
    store, end = find_store(fd)
    first = (store + 28 + 3) & ~3
    if struct.unpack_from("<H", fd, first)[0] != 0xFFFF:
        fail("the firmware's variable store is not empty; seed only a release RPI_EFI.fd")

    blob = (variable("Boot0000", load_option("MontaukOS", disk_id, part, start, sectors))
            + variable("BootOrder", struct.pack("<H", 0)))
    if first + len(blob) > end:
        fail("the variable store is too small")
    fd[first:first + len(blob)] = blob
    open(out, "wb").write(fd)


if __name__ == "__main__":
    main()