microdroid/extract_microdroid_kernel_hashes.py

"""Extracts the following hashes from the AVB footer of Microdroid's kernel:

- kernel hash
- initrd_normal hash
- initrd_debug hash

The hashes are written to stdout as a Rust file.

In unsupportive environments such as x86, when the kernel is just an empty file,
the output Rust file has the same hash constant fields for compatibility
reasons, but all of them are empty.
"""
#!/usr/bin/env python3

import argparse
from collections import defaultdict
import subprocess
from typing import Dict

PARTITION_NAME_BOOT = 'boot'
PARTITION_NAME_INITRD_NORMAL = 'initrd_normal'
PARTITION_NAME_INITRD_DEBUG = 'initrd_debug'
HASH_SIZE = 32

def main(args):
    """Main function."""
    avbtool = args.avbtool
    num_kernel_images = len(args.kernel)

    print("//! This file is generated by extract_microdroid_kernel_hashes.py.")
    print("//! It contains the hashes of the kernel and initrds.\n")
    print("#![no_std]\n#![allow(missing_docs)]\n")

    print("pub const HASH_SIZE: usize = " + str(HASH_SIZE) + ";\n")
    print("pub struct OsHashes {")
    print("    pub kernel: [u8; HASH_SIZE],")
    print("    pub initrd_normal: [u8; HASH_SIZE],")
    print("    pub initrd_debug: [u8; HASH_SIZE],")
    print("}\n")

    hashes = defaultdict(list)
    for kernel_image_path in args.kernel:
        collected_hashes = collect_hashes(avbtool, kernel_image_path)

        if collected_hashes.keys() == {PARTITION_NAME_BOOT,
                                       PARTITION_NAME_INITRD_NORMAL,
                                       PARTITION_NAME_INITRD_DEBUG}:
            for partition_name, v in collected_hashes.items():
                hashes[partition_name].append(v)
        else:
            # Microdroid's kernel is just an empty file in unsupportive
            # environments such as x86, in this case the hashes should be empty.
            print("/// The kernel is empty, no hashes are available.")
            hashes[PARTITION_NAME_BOOT].append("")
            hashes[PARTITION_NAME_INITRD_NORMAL].append("")
            hashes[PARTITION_NAME_INITRD_DEBUG].append("")

    print("pub const OS_HASHES: [OsHashes; " + str(num_kernel_images) + "] = [")
    for i in range(num_kernel_images):
        print("OsHashes {")
        print("    kernel: [" +
              format_hex_string(hashes[PARTITION_NAME_BOOT][i]) + "],")
        print("    initrd_normal: [" +
              format_hex_string(hashes[PARTITION_NAME_INITRD_NORMAL][i]) + "],")
        print("    initrd_debug: [" +
              format_hex_string(hashes[PARTITION_NAME_INITRD_DEBUG][i]) + "],")
        print("},")
    print("];")

def collect_hashes(avbtool: str, kernel_image_path: str) -> Dict[str, str]:
    """Collects the hashes from the AVB footer of the kernel image."""
    hashes = {}
    with subprocess.Popen(
        [avbtool, 'print_partition_digests', '--image', kernel_image_path],
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT) as proc:
        stdout, _ = proc.communicate()
        for line in stdout.decode("utf-8").split("\n"):
            line = line.replace(" ", "").split(":")
            if len(line) == 2:
                partition_name, hash_ = line
                hashes[partition_name] = hash_
    return hashes

def format_hex_string(hex_string: str) -> str:
    """Formats a hex string into a Rust array."""
    if not hex_string:
        return "0x00, " * HASH_SIZE
    assert len(hex_string) == HASH_SIZE * 2, \
          "Hex string must have length " + str(HASH_SIZE * 2) + ": " + \
          hex_string
    return ", ".join(["\n0x" + hex_string[i:i+2] if i % 32 == 0
                       else "0x" + hex_string[i:i+2]
                       for i in range(0, len(hex_string), 2)])

def parse_args():
    """Parses the command line arguments."""
    parser = argparse.ArgumentParser(
        "Extracts the hashes from the kernels' AVB footer")
    parser.add_argument('--avbtool', help='Path to the avbtool binary')
    parser.add_argument('--kernel', help='Path to the kernel image', nargs='+')
    return parser.parse_args()

if __name__ == '__main__':
    main(parse_args())