Add new bootimagetool utility for working with boot images

This just adds a simple CLI for avbroot.formats.bootimage library. It's
not strictly related to avbroot, but is useful for troubleshooting since
the library supports all AOSP-compatible boot images.

Signed-off-by: Andrew Gunnerson <accounts+github@chiller3.com>

avbroot/formats/bootimage.py

@@ -113,7 +113,13 @@ class WrongFormat(ValueError):
 
 
 class BootImage:
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        f: None | typing.BinaryIO = None,
+        data: None | dict[str, typing.Any] = None,
+    ) -> None:
+        assert (f is None) != (data is None)
+
         self.kernel: None | bytes = None
         self.ramdisks: list[bytes] = []
         self.second: None | bytes = None
@@ -121,14 +127,26 @@ class BootImage:
         self.dtb: None | bytes = None
         self.bootconfig: None | bytes = None
 
+        if f:
+            self._from_file(f)
+        else:
+            self._from_dict(data)
+
+    def _from_file(self, f: typing.BinaryIO) -> None:
+        raise NotImplementedError()
+
     def generate(self, f: typing.BinaryIO) -> None:
         raise NotImplementedError()
 
+    def _from_dict(self, data: dict[str, typing.Any]) -> None:
+        raise NotImplementedError()
+
+    def to_dict(self) -> None:
+        raise NotImplementedError()
+
 
 class _BootImageV0Through2(BootImage):
-    def __init__(self, f: typing.BinaryIO) -> None:
-        super().__init__()
-
+    def _from_file(self, f: typing.BinaryIO) -> None:
         # Common fields for v0 through v2
         magic, kernel_size, kernel_addr, ramdisk_size, ramdisk_addr, \
             second_size, second_addr, tags_addr, page_size, header_version, \
@@ -299,11 +317,60 @@ class _BootImageV0Through2(BootImage):
 
         return result
 
+    def _from_dict(self, data: dict[str, typing.Any]) -> None:
+        type = data.get('type')
+        header_version = data.get('header_version')
+
+        if type != 'android':
+            raise WrongFormat(f'Unknown type: {type}')
+        elif header_version not in (0, 1, 2):
+            raise WrongFormat(f'Unknown header version: {header_version}')
+
+        self.header_version = header_version
+        self.kernel_addr = data['kernel_address']
+        self.ramdisk_addr = data['ramdisk_address']
+        self.second_addr = data['second_address']
+        self.tags_addr = data['tags_address']
+        self.page_size = data['page_size']
+        self.os_version = data['os_version']
+        self.name = data['name']
+        self.cmdline = data['cmdline']
+        self.id = data['id']
+        self.extra_cmdline = data['extra_cmdline']
+
+        if header_version >= 1:
+            self.recovery_dtbo_offset = data['recovery_dtbo_offset']
+
+        if self.header_version == 2:
+            self.dtb_addr = data['dtb_address']
+
+    def to_dict(self) -> dict[str, typing.Any]:
+        result = {
+            'type': 'android',
+            'header_version': self.header_version,
+            'kernel_address': self.kernel_addr,
+            'ramdisk_address': self.ramdisk_addr,
+            'second_address': self.second_addr,
+            'tags_address': self.tags_addr,
+            'page_size': self.page_size,
+            'os_version': self.os_version,
+            'name': self.name,
+            'cmdline': self.cmdline,
+            'id': self.id,
+            'extra_cmdline': self.extra_cmdline,
+        }
+
+        if self.header_version >= 1:
+            result['recovery_dtbo_offset'] = self.recovery_dtbo_offset
+
+        if self.header_version == 2:
+            result['dtb_address'] = self.dtb_addr
+
+        return result
+
 
 class _BootImageV3Through4(BootImage):
-    def __init__(self, f: typing.BinaryIO) -> None:
-        super().__init__()
-
+    def _from_file(self, f: typing.BinaryIO) -> None:
         # Common fields for both v3 and v4
         magic, kernel_size, ramdisk_size, os_version, header_size, reserved, \
             header_version, cmdline = BOOT_IMG_HDR_V3.unpack(
@@ -395,15 +462,36 @@ class _BootImageV3Through4(BootImage):
             f'- Reserved:       {self.reserved.hex()}\n' \
             f'- Kernel cmdline: {self.cmdline!r}\n'
 
+    def _from_dict(self, data: dict[str, typing.Any]) -> None:
+        type = data.get('type')
+        header_version = data.get('header_version')
+
+        if type != 'android':
+            raise WrongFormat(f'Unknown type: {type}')
+        elif header_version not in (3, 4):
+            raise WrongFormat(f'Unknown header version: {header_version}')
+
+        self.header_version = header_version
+        self.os_version = data['os_version']
+        self.reserved = data['reserved']
+        self.cmdline = data['cmdline']
+
+    def to_dict(self) -> dict[str, typing.Any]:
+        return {
+            'type': 'android',
+            'header_version': self.header_version,
+            'os_version': self.os_version,
+            'reserved': self.reserved,
+            'cmdline': self.cmdline,
+        }
+
 
 _RamdiskMeta = collections.namedtuple(
     '_RamdiskMeta', ['type', 'name', 'board_id'])
 
 
 class _VendorBootImageV3Through4(BootImage):
-    def __init__(self, f: typing.BinaryIO) -> None:
-        super().__init__()
-
+    def _from_file(self, f: typing.BinaryIO) -> None:
         # Common fields for both v3 and v4
         magic, header_version, page_size, kernel_addr, ramdisk_addr, \
             vendor_ramdisk_size, cmdline, tags_addr, name, header_size, \
@@ -573,7 +661,13 @@ class _VendorBootImageV3Through4(BootImage):
         result = \
             f'Vendor boot image v{self.header_version} header:\n' \
             f'- Page size:           {self.page_size}\n' \
-            f'- Kernel address:      0x{self.kernel_addr:x}\n' \
+            f'- Kernel address:      0x{self.kernel_addr:x}\n'
+
+        if self.header_version == 3:
+            ramdisk_size = len(self.ramdisks[0]) if self.ramdisks else 0
+            result += f'- Ramdisk size:        {ramdisk_size}\n'
+
+        result += \
             f'- Ramdisk address:     0x{self.ramdisk_addr:x}\n' \
             f'- Kernel cmdline:      {self.cmdline!r}\n' \
             f'- Kernel tags address: 0x{self.tags_addr:x}\n' \
@@ -596,6 +690,57 @@ class _VendorBootImageV3Through4(BootImage):
 
         return result
 
+    def _from_dict(self, data: dict[str, typing.Any]) -> None:
+        type = data.get('type')
+        header_version = data.get('header_version')
+
+        if type != 'vendor':
+            raise WrongFormat(f'Unknown type: {type}')
+        elif header_version not in (3, 4):
+            raise WrongFormat(f'Unknown header version: {header_version}')
+
+        self.header_version = header_version
+        self.page_size = data['page_size']
+        self.kernel_addr = data['kernel_address']
+        self.ramdisk_addr = data['ramdisk_address']
+        self.cmdline = data['cmdline']
+        self.tags_addr = data['tags_address']
+        self.name = data['name']
+        self.dtb_addr = data['dtb_address']
+
+        if header_version == 4:
+            self.ramdisks_meta = []
+            for meta in data['ramdisk_meta']:
+                self.ramdisks_meta.append(_RamdiskMeta(
+                    meta['type'],
+                    meta['name'],
+                    meta['board_id'],
+                ))
+
+    def to_dict(self) -> dict[str, typing.Any]:
+        result = {
+            'type': 'vendor',
+            'header_version': self.header_version,
+            'page_size': self.page_size,
+            'kernel_address': self.kernel_addr,
+            'ramdisk_address': self.ramdisk_addr,
+            'cmdline': self.cmdline,
+            'tags_address': self.tags_addr,
+            'name': self.name,
+            'dtb_address': self.dtb_addr,
+        }
+
+        if self.header_version == 4:
+            result['ramdisk_meta'] = []
+            for meta in self.ramdisks_meta:
+                result['ramdisk_meta'].append({
+                    'type': meta.type,
+                    'name': meta.name,
+                    'board_id': meta.board_id,
+                })
+
+        return result
+
 
 def load_autodetect(f: typing.BinaryIO) -> BootImage:
     for cls in (
@@ -605,7 +750,21 @@ def load_autodetect(f: typing.BinaryIO) -> BootImage:
     ):
         try:
             f.seek(0)
-            return cls(f)
+            return cls(f=f)
+        except WrongFormat:
+            continue
+
+    raise ValueError('Unknown boot image format')
+
+
+def create_from_dict(data: dict) -> BootImage:
+    for cls in (
+        _BootImageV0Through2,
+        _BootImageV3Through4,
+        _VendorBootImageV3Through4,
+    ):
+        try:
+            return cls(data=data)
         except WrongFormat:
             continue
 

extra/README.md

@@ -0,0 +1,31 @@
+# avbroot extra
+
+This directory contains some extra scripts that aren't required for avbroot's operation, but may be useful for troubleshooting.
+
+## `bootimagetool`
+
+This is a frontend to the [`avbroot/formats/bootimage.py`](../avbroot/formats/bootimage.py) library for working with boot images.
+
+### Unpacking a boot image
+
+```bash
+python bootimagetool.py unpack <input boot image>
+```
+
+This subcommand unpacks all of the components of the boot image into the current directory by default (see `--help`). The header fields are saved to `header.json` and each blob section is saved to a separate file. Each blob is written to disk as-is, without decompression.
+
+### Packing a boot image
+
+```bash
+python bootimagetool.py pack <output boot image>
+```
+
+This subcommand packs a new boot image from the individual components in the current directory by default (see `--help`). The default input filenames are the same as the output filenames for the `unpack` subcommand.
+
+### Repacking a boot image
+
+```bash
+python bootimagetool.py repack <input boot image> <output boot image>
+```
+
+This subcommand repacks a boot image without writing the individual components to disk first. This is useful for roundtrip testing of avbroot's boot image parser. The output should be identical to the input, minus any footers, like the AVB footer. The only exception is the VTS signature for v4 boot images, which is always stripped out.

extra/bootimagetool.py

@@ -0,0 +1,168 @@
+#!/usr/bin/env python3
+
+import argparse
+import itertools
+import json
+import os
+import sys
+
+sys.path.append(os.path.join(sys.path[0], '..'))
+from avbroot.formats import bootimage
+
+
+class BytesDecoder(json.JSONDecoder):
+    def __init__(self):
+        super().__init__(object_hook=self.from_dict)
+
+    @staticmethod
+    def from_dict(d):
+        # This is insufficient for arbitrary data, but we're not dealing with
+        # arbitrary data
+        if 'type' in d:
+            if d['type'] == 'UTF-8':
+                return d['data'].encode('UTF-8')
+            elif d['type'] == 'hex':
+                return bytes.fromhex(d['data'])
+
+        return d
+
+
+class BytesEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, bytes):
+            if b'\0' not in obj:
+                try:
+                    return {
+                        'type': 'UTF-8',
+                        'data': obj.decode('UTF-8'),
+                    }
+                except UnicodeDecodeError:
+                    pass
+
+            return {
+                'type': 'hex',
+                'data': obj.hex(),
+            }
+
+        return super().default(obj)
+
+
+def read_or_none(path):
+    try:
+        with open(path, 'rb') as f:
+            return f.read()
+    except FileNotFoundError:
+        return None
+
+
+def write_if_not_none(path, data):
+    if data is not None:
+        with open(path, 'wb') as f:
+            f.write(data)
+
+
+def parse_args():
+    parser_kwargs = {'formatter_class': argparse.ArgumentDefaultsHelpFormatter}
+
+    parser = argparse.ArgumentParser(**parser_kwargs)
+    subparsers = parser.add_subparsers(dest='subcommand', required=True,
+                                       help='Subcommands')
+
+    base = argparse.ArgumentParser(add_help=False)
+    base.add_argument('-q', '--quiet', action='store_true',
+                      help='Do not print header information')
+
+    pack = subparsers.add_parser('pack', help='Pack a boot image',
+                                 parents=[base], **parser_kwargs)
+    unpack = subparsers.add_parser('unpack', help='Unpack a boot image',
+                                   parents=[base], **parser_kwargs)
+    repack = subparsers.add_parser('repack', help='Repack a boot image',
+                                   parents=[base], **parser_kwargs)
+
+    for p in (pack, unpack):
+        prefix = '--input-' if p == pack else '--output-'
+
+        p.add_argument('boot_image', help='Path to boot image')
+
+        p.add_argument(prefix + 'header', default='header.json',
+                       help='Path to header JSON')
+        p.add_argument(prefix + 'kernel', default='kernel.img',
+                       help='Path to kernel')
+        p.add_argument(prefix + 'ramdisk-prefix', default='ramdisk.img.',
+                       help='Path prefix for ramdisk')
+        p.add_argument(prefix + 'second', default='second.img',
+                       help='Path to second stage bootloader')
+        p.add_argument(prefix + 'recovery-dtbo', default='recovery_dtbo.img',
+                       help='Path to recovery dtbo/acpio')
+        p.add_argument(prefix + 'dtb', default='dtb.img',
+                       help='Path to device tree blob')
+        p.add_argument(prefix + 'bootconfig', default='bootconfig.txt',
+                       help='Path to bootconfig')
+
+    repack.add_argument('input', help='Path to input boot image')
+    repack.add_argument('output', help='Path to output boot image')
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    if args.subcommand == 'pack':
+        with open(args.input_header, 'r') as f:
+            data = json.load(f, cls=BytesDecoder)
+
+        img = bootimage.create_from_dict(data)
+
+        img.kernel = read_or_none(args.input_kernel)
+        img.second = read_or_none(args.input_second)
+        img.recovery_dtbo = read_or_none(args.input_recovery_dtbo)
+        img.dtb = read_or_none(args.input_dtb)
+        img.bootconfig = read_or_none(args.input_bootconfig)
+
+        for i in itertools.count():
+            ramdisk = read_or_none(f'{args.input_ramdisk_prefix}{i}')
+            if ramdisk is None:
+                break
+
+            img.ramdisks.append(ramdisk)
+
+        if not args.quiet:
+            print(img)
+
+        with open(args.boot_image, 'wb') as f:
+            img.generate(f)
+
+    elif args.subcommand == 'unpack':
+        with open(args.boot_image, 'rb') as f:
+            img = bootimage.load_autodetect(f)
+            if not args.quiet:
+                print(img)
+
+        with open(args.output_header, 'w') as f:
+            json.dump(img.to_dict(), f, indent=4, cls=BytesEncoder)
+
+        write_if_not_none(args.output_kernel, img.kernel)
+        write_if_not_none(args.output_second, img.second)
+        write_if_not_none(args.output_recovery_dtbo, img.recovery_dtbo)
+        write_if_not_none(args.output_dtb, img.dtb)
+        write_if_not_none(args.output_bootconfig, img.bootconfig)
+
+        for i, ramdisk in enumerate(img.ramdisks):
+            write_if_not_none(f'{args.output_ramdisk_prefix}{i}', ramdisk)
+
+    elif args.subcommand == 'repack':
+        with open(args.input, 'rb') as f:
+            img = bootimage.load_autodetect(f)
+            if not args.quiet:
+                print(img)
+
+        with open(args.output, 'wb') as f:
+            img.generate(f)
+
+    else:
+        raise NotImplementedError()
+
+
+if __name__ == '__main__':
+    main()

tests/tests.conf

@@ -8,7 +8,7 @@ url = https://dl.google.com/dl/android/aosp/cheetah-ota-tq1a.230205.002-ad27a6d7
 sha256 = ad27a6d7ca78c0d3fb7f51864cc6a9e6bc91f536f9e9bab5da06253df602525f
 sha256_patched = a711b2033a6ebe39c2ba22e885bd570db33092897eda1793e508f3e8ffb8000a
 
-# What's unique: boot (boot v4, no ramdisk) + vendor_boot (vendor v4)
+# What's unique: boot (boot v4, no ramdisk) + vendor_boot (vendor v4, 2 ramdisks)
 [device.GooglePixel6a]
 url = https://dl.google.com/dl/android/aosp/bluejay-ota-tq1a.230205.002-1ec64b58.zip
 sha256 = 1ec64b589cba8d9687ea8f85cf6c2e5d09c9e4fbdae15417aa06d5ba976fbe79