pack-uboot-boot.py

#!/usr/bin/env python3
"""Pack u-boot.bin into an Android boot image that MediaTek lk will accept.

Deliberately does not use mkbootimg and does not invent load addresses: it takes
the header of a boot.img that is known to boot on this phone and swaps only the
kernel area. That way every address (kernel/ramdisk/tags/dtb) is byte-identical
to the image lk already boots.

Two traps, both of which make lk die before printing anything:
  - ramdisk_size = 0  -> lk_crash. The template's ramdisk must be kept.
  - no AVB footer     -> lk_crash. Run avbtool after packing.

  ./pack-uboot-boot.py <template.img> <payload> <out.img> [replacement-dtb]
"""
import gzip
import hashlib
import io
import struct
import sys

MAGIC = b"ANDROID!"


def roundup(n, page):
    return (n + page - 1) // page * page


def main(template, uboot, out, new_dtb=None):
    d = bytearray(open(template, "rb").read())
    if bytes(d[:8]) != MAGIC:
        sys.exit(f"{template}: not an Android boot image")

    (kernel_size, kernel_addr, ramdisk_size, ramdisk_addr,
     second_size, second_addr, tags_addr, page_size,
     header_version) = struct.unpack_from("<9I", d, 8)

    dtb_size, = struct.unpack_from("<I", d, 1648)

    print(f"template    : {template}")
    print(f"  page_size {page_size}  header v{header_version}")
    print(f"  kernel  {kernel_size:>9} B @ {kernel_addr:#010x}")
    print(f"  ramdisk {ramdisk_size:>9} B @ {ramdisk_addr:#010x}")
    print(f"  second  {second_size:>9} B")
    print(f"  dtb     {dtb_size:>9} B")

    if ramdisk_size == 0:
        sys.exit("template has no ramdisk: lk will report lk_crash")

    hdr_pages = roundup(1648 + 12, page_size)
    off = hdr_pages
    kernel_off = off
    off += roundup(kernel_size, page_size)
    ramdisk_off = off
    off += roundup(ramdisk_size, page_size)
    second_off = off
    off += roundup(second_size, page_size)
    dtb_off = off

    ramdisk = bytes(d[ramdisk_off:ramdisk_off + ramdisk_size])
    second = bytes(d[second_off:second_off + second_size])
    dtb = bytes(d[dtb_off:dtb_off + dtb_size])
    if new_dtb:
        dtb = open(new_dtb, "rb").read()
        print(f"replacement dtb: {new_dtb}  {len(dtb)} B")

    payload = open(uboot, "rb").read()
    if payload[56:60] != b"ARM\x64":
        print("  warning: no ARM64 magic at offset 56 of u-boot.bin;"
              " CONFIG_LINUX_KERNEL_IMAGE_HEADER=y is required")
    print(f"u-boot      : {uboot}  {len(payload)} B")

    # lk on this phone GUNZIPS the kernel slot before jumping into it. The
    # template's kernel starts with 1f 8b 08, i.e. gzip. Putting a raw ARM64
    # Image there makes lk's decompressor fault and the phone comes back with
    # androidboot.bootreason=lk_crash, having executed nothing.
    kernel_orig = bytes(d[kernel_off:kernel_off + kernel_size])
    if kernel_orig[:3] == b"\x1f\x8b\x08":
        buf = io.BytesIO()
        # mtime=0 so the output is reproducible
        with gzip.GzipFile(fileobj=buf, mode="wb", compresslevel=9, mtime=0) as g:
            g.write(payload)
        new_kernel = buf.getvalue()
        print(f"  template kernel is gzip -> recompressing u-boot: {len(new_kernel)} B")
    else:
        new_kernel = payload
        print("  template kernel is raw -> keeping payload as is")

    head = bytearray(d[:hdr_pages])
    struct.pack_into("<I", head, 8, len(new_kernel))
    struct.pack_into("<I", head, 1648, len(dtb))

    img = bytearray()
    img += head
    img += new_kernel + b"\0" * (roundup(len(new_kernel), page_size) - len(new_kernel))
    img += ramdisk + b"\0" * (roundup(len(ramdisk), page_size) - len(ramdisk))
    if second:
        img += second + b"\0" * (roundup(len(second), page_size) - len(second))
    if dtb:
        img += dtb + b"\0" * (roundup(len(dtb), page_size) - len(dtb))

    # The id field is a SHA1 over each section followed by its length. Keeping
    # the template's id would leave it inconsistent with the new contents.
    sha = hashlib.sha1()
    for part in (new_kernel, ramdisk, second):
        sha.update(part)
        sha.update(struct.pack("<I", len(part)))
    if dtb:
        sha.update(dtb)
        sha.update(struct.pack("<I", len(dtb)))
    digest = sha.digest()
    img[576:576 + 32] = digest + b"\0" * (32 - len(digest))
    print(f"  id (SHA1) recomputed: {digest.hex()}")

    open(out, "wb").write(bytes(img))
    print(f"result      : {out}  {len(img)} B")
    print()
    print("Next step (mandatory, lk_crash without it):")
    print(f"  python3 avbtool.py add_hash_footer --image {out} \\")
    print("      --partition_name recovery --partition_size 67108864 --algorithm NONE")


if __name__ == "__main__":
    if len(sys.argv) not in (4, 5):
        sys.exit(__doc__)
    main(*sys.argv[1:])