1"""Extracts the following hashes from the AVB footer of Microdroid's kernel:
2
3- kernel hash
4- initrd_normal hash
5- initrd_debug hash
6
7The hashes are written to stdout as a Rust file.
8
9In unsupportive environments such as x86, when the kernel is just an empty file,
10the output Rust file has the same hash constant fields for compatibility
11reasons, but all of them are empty.
12"""
13#!/usr/bin/env python3
14
15import argparse
16from collections import defaultdict
17import subprocess
18from typing import Dict
19
20PARTITION_NAME_BOOT = 'boot'
21PARTITION_NAME_INITRD_NORMAL = 'initrd_normal'
22PARTITION_NAME_INITRD_DEBUG = 'initrd_debug'
23HASH_SIZE = 32
24
25def main(args):
26    """Main function."""
27    avbtool = args.avbtool
28    num_kernel_images = len(args.kernel)
29
30    print("//! This file is generated by extract_microdroid_kernel_hashes.py.")
31    print("//! It contains the hashes of the kernel and initrds.\n")
32    print("#![no_std]\n#![allow(missing_docs)]\n")
33
34    print("pub const HASH_SIZE: usize = " + str(HASH_SIZE) + ";\n")
35    print("pub struct OsHashes {")
36    print("    pub kernel: [u8; HASH_SIZE],")
37    print("    pub initrd_normal: [u8; HASH_SIZE],")
38    print("    pub initrd_debug: [u8; HASH_SIZE],")
39    print("}\n")
40
41    hashes = defaultdict(list)
42    for kernel_image_path in args.kernel:
43        collected_hashes = collect_hashes(avbtool, kernel_image_path)
44
45        if collected_hashes.keys() == {PARTITION_NAME_BOOT,
46                                       PARTITION_NAME_INITRD_NORMAL,
47                                       PARTITION_NAME_INITRD_DEBUG}:
48            for partition_name, v in collected_hashes.items():
49                hashes[partition_name].append(v)
50        else:
51            # Microdroid's kernel is just an empty file in unsupportive
52            # environments such as x86, in this case the hashes should be empty.
53            print("/// The kernel is empty, no hashes are available.")
54            hashes[PARTITION_NAME_BOOT].append("")
55            hashes[PARTITION_NAME_INITRD_NORMAL].append("")
56            hashes[PARTITION_NAME_INITRD_DEBUG].append("")
57
58    print("pub const OS_HASHES: [OsHashes; " + str(num_kernel_images) + "] = [")
59    for i in range(num_kernel_images):
60        print("OsHashes {")
61        print("    kernel: [" +
62              format_hex_string(hashes[PARTITION_NAME_BOOT][i]) + "],")
63        print("    initrd_normal: [" +
64              format_hex_string(hashes[PARTITION_NAME_INITRD_NORMAL][i]) + "],")
65        print("    initrd_debug: [" +
66              format_hex_string(hashes[PARTITION_NAME_INITRD_DEBUG][i]) + "],")
67        print("},")
68    print("];")
69
70def collect_hashes(avbtool: str, kernel_image_path: str) -> Dict[str, str]:
71    """Collects the hashes from the AVB footer of the kernel image."""
72    hashes = {}
73    with subprocess.Popen(
74        [avbtool, 'print_partition_digests', '--image', kernel_image_path],
75        stdout=subprocess.PIPE, stderr=subprocess.STDOUT) as proc:
76        stdout, _ = proc.communicate()
77        for line in stdout.decode("utf-8").split("\n"):
78            line = line.replace(" ", "").split(":")
79            if len(line) == 2:
80                partition_name, hash_ = line
81                hashes[partition_name] = hash_
82    return hashes
83
84def format_hex_string(hex_string: str) -> str:
85    """Formats a hex string into a Rust array."""
86    if not hex_string:
87        return "0x00, " * HASH_SIZE
88    assert len(hex_string) == HASH_SIZE * 2, \
89          "Hex string must have length " + str(HASH_SIZE * 2) + ": " + \
90          hex_string
91    return ", ".join(["\n0x" + hex_string[i:i+2] if i % 32 == 0
92                       else "0x" + hex_string[i:i+2]
93                       for i in range(0, len(hex_string), 2)])
94
95def parse_args():
96    """Parses the command line arguments."""
97    parser = argparse.ArgumentParser(
98        "Extracts the hashes from the kernels' AVB footer")
99    parser.add_argument('--avbtool', help='Path to the avbtool binary')
100    parser.add_argument('--kernel', help='Path to the kernel image', nargs='+')
101    return parser.parse_args()
102
103if __name__ == '__main__':
104    main(parse_args())
105