Move kernel hash extraction scripts to microdroid

Test: atest rialto_test
Bug: 326363997
Change-Id: Iecc53ef6cfecc8b2a89290b1d469a2940000c27e
diff --git a/microdroid/Android.bp b/microdroid/Android.bp
index 233754a..5940835 100644
--- a/microdroid/Android.bp
+++ b/microdroid/Android.bp
@@ -599,3 +599,37 @@
     defaults: ["microdroid_initrd_debug_defaults"],
     src: ":microdroid_gki-android14-6.1_initrd_debuggable",
 }
+
+python_binary_host {
+    name: "extract_microdroid_kernel_hashes",
+    srcs: ["extract_microdroid_kernel_hashes.py"],
+}
+
+genrule {
+    name: "microdroid_kernel_hashes_rs",
+    srcs: [
+        ":microdroid_kernel",
+        ":microdroid_gki-android14-6.1_kernel",
+    ],
+    out: ["lib.rs"],
+    tools: [
+        "extract_microdroid_kernel_hashes",
+        "avbtool",
+    ],
+    cmd: "$(location extract_microdroid_kernel_hashes) --avbtool $(location avbtool) " +
+        "--kernel $(location :microdroid_kernel) " +
+        "$(location :microdroid_gki-android14-6.1_kernel) " +
+        "> $(out)",
+}
+
+rust_library_rlib {
+    name: "libmicrodroid_kernel_hashes",
+    srcs: [":microdroid_kernel_hashes_rs"],
+    crate_name: "microdroid_kernel_hashes",
+    prefer_rlib: true,
+    no_stdlibs: true,
+    stdlibs: [
+        "libcompiler_builtins.rust_sysroot",
+        "libcore.rust_sysroot",
+    ],
+}
diff --git a/microdroid/extract_microdroid_kernel_hashes.py b/microdroid/extract_microdroid_kernel_hashes.py
new file mode 100644
index 0000000..f2c6ae7
--- /dev/null
+++ b/microdroid/extract_microdroid_kernel_hashes.py
@@ -0,0 +1,104 @@
+"""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())