Snap for 8570526 from 0674b7ccb3dc02466e0e08da28693d04b56f6a7d to mainline-mediaprovider-release

Change-Id: I13f286383b8f72418ad7528bbe29cf548a4f9db7
diff --git a/Android.bp b/Android.bp
index 23c55b8..65b6ac2 100644
--- a/Android.bp
+++ b/Android.bp
@@ -56,8 +56,10 @@
 python_binary_host {
     name: "mkbootimg",
     defaults: ["mkbootimg_defaults"],
+    main: "mkbootimg.py",
     srcs: [
         "mkbootimg.py",
+        "gki/generate_gki_certificate.py",
     ],
     required: [
         "avbtool",
@@ -89,6 +91,20 @@
     ],
 }
 
+python_binary_host {
+    name: "certify_bootimg",
+    defaults: ["mkbootimg_defaults"],
+    main: "gki/certify_bootimg.py",
+    srcs: [
+        "gki/certify_bootimg.py",
+        "gki/generate_gki_certificate.py",
+        "unpack_bootimg.py",
+    ],
+    required: [
+        "avbtool",
+    ],
+}
+
 python_test_host {
     name: "mkbootimg_test",
     defaults: ["mkbootimg_defaults"],
diff --git a/BUILD.bazel b/BUILD.bazel
new file mode 100644
index 0000000..e80a82a
--- /dev/null
+++ b/BUILD.bazel
@@ -0,0 +1,17 @@
+# Copyright (C) 2021 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+exports_files([
+    "mkbootimg.py",
+])
diff --git a/OWNERS b/OWNERS
index 51e09a2..d71a5a1 100644
--- a/OWNERS
+++ b/OWNERS
@@ -1,3 +1,2 @@
-hridya@google.com
 smuckle@google.com
 yochiang@google.com
\ No newline at end of file
diff --git a/gki/Android.bp b/gki/Android.bp
new file mode 100644
index 0000000..c62e7d8
--- /dev/null
+++ b/gki/Android.bp
@@ -0,0 +1,103 @@
+// Copyright (C) 2022 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+python_test_host {
+    name: "certify_bootimg_test",
+    defaults: ["mkbootimg_defaults"],
+    main: "certify_bootimg_test.py",
+    srcs: [
+        "certify_bootimg_test.py",
+    ],
+    data: [
+        ":avbtool",
+        ":certify_bootimg",
+        ":mkbootimg",
+        ":unpack_bootimg",
+        "testdata/*",
+    ],
+    test_options: {
+        unit_test: true,
+    },
+}
+
+python_binary_host {
+    name: "generate_gki_certificate",
+    defaults: ["mkbootimg_defaults"],
+    srcs: [
+        "generate_gki_certificate.py",
+    ],
+    required: [
+        "avbtool",
+    ],
+}
+
+sh_binary_host {
+    name: "retrofit_gki",
+    src: "retrofit_gki.sh",
+    required: [
+        "avbtool",
+        "mkbootimg",
+        "unpack_bootimg",
+    ],
+}
+
+sh_test_host {
+    name: "retrofit_gki_test",
+    src: "retrofit_gki_test.sh",
+    data: [
+        "retrofit_gki.sh",
+    ],
+    data_bins: [
+        "avbtool",
+        "mkbootimg",
+        "unpack_bootimg",
+    ],
+    test_suites: [
+        "general-tests",
+    ],
+}
+
+genrule {
+    name: "gki_retrofitting_tools",
+    tools: [
+        "soong_zip",
+        "retrofit_gki",
+        "avbtool",
+        "mkbootimg",
+        "unpack_bootimg",
+    ],
+    srcs: [
+        "README.md",
+    ],
+    cmd: "STAGE_DIR=$(genDir)/gki_retrofitting_tools && " +
+         "rm -rf $${STAGE_DIR} && mkdir -p $${STAGE_DIR} && " +
+         "cp $(location retrofit_gki) $${STAGE_DIR} && " +
+         "cp $(location avbtool) $${STAGE_DIR} && " +
+         "cp $(location mkbootimg) $${STAGE_DIR} && " +
+         "cp $(location unpack_bootimg) $${STAGE_DIR} && " +
+         "cp $(in) $${STAGE_DIR} && " +
+         "$(location soong_zip) -o $(out) -C $(genDir) -D $${STAGE_DIR}",
+    out: [
+        "gki_retrofitting_tools.zip",
+    ],
+    dist: {
+        targets: [
+            "gki_retrofitting_tools",
+        ],
+    },
+}
diff --git a/gki/Android.mk b/gki/Android.mk
new file mode 100644
index 0000000..c0af5ef
--- /dev/null
+++ b/gki/Android.mk
@@ -0,0 +1,36 @@
+#
+# Copyright (C) 2022 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+_gsi_gki_product_names := \
+  aosp_arm \
+  aosp_arm64 \
+  aosp_x86 \
+  aosp_x86_64 \
+  gsi_arm \
+  gsi_arm64 \
+  gsi_x86 \
+  gsi_x86_64 \
+  gki_arm64 \
+  gki_x86_64 \
+
+# Add gki_retrofitting_tools to `m dist` of GSI/GKI for easy pickup.
+ifneq (,$(filter $(_gsi_gki_product_names),$(TARGET_PRODUCT)))
+
+droidcore-unbundled: gki_retrofitting_tools
+
+endif
+
+_gsi_gki_product_names :=
diff --git a/gki/README.md b/gki/README.md
new file mode 100644
index 0000000..628e52a
--- /dev/null
+++ b/gki/README.md
@@ -0,0 +1,74 @@
+# GKI boot image retrofitting tools for upgrading devices
+
+Starting from Android T the GKI boot images consist of the generic `boot.img`
+and `init_boot.img`. The `boot.img` contains the generic kernel, and
+`init_boot.img` contains the generic ramdisk.
+For upgrading devices whose `vendor_boot` partition is non-existent, this tool
+(or spec) can be used to retrofit a set of Android T GKI `boot`, `init_boot` and
+OEM `vendor_boot` partition images back into a single boot image containing the
+GKI kernel plus generic and vendor ramdisks.
+
+## Retrofitting the boot images
+
+1. Download the certified GKI `boot.img`.
+2. Go to the build artifacts page of `aosp_arm64` on `aosp-master` branch on
+   https://ci.android.com/ and download `gki_retrofitting_tools.zip`.
+3. Unzip and make sure the tool is in `${PATH}`.
+
+   ```bash
+   unzip gki_retrofitting_tools.zip
+   export PATH="$(pwd)/gki_retrofitting_tools:${PATH}"
+   # See tool usage:
+   retrofit_gki --help
+   ```
+
+4. Create the retrofitted image. The `--version` argument lets you choose the
+   boot image header version of the retrofitted boot image. Only version 2 is
+   supported at the moment.
+
+   ```bash
+   retrofit_gki --boot boot.img --init_boot init_boot.img \
+     --vendor_boot vendor_boot.img --version 2 -o boot.retrofitted.img
+   ```
+
+## Spec of the retrofitted images
+
+* The SOURCE `boot.img` must be officially certified Android T (or later) GKI.
+* The DEST retrofitted boot image must not set the security patch level in its
+  header. This is because the SOURCE images might have different SPL value, thus
+  making the boot header SPL of the retrofitted image ill-defined. The SPL value
+  must be defined by the chained vbmeta image of the `boot` partition.
+* The `boot signature` of the DEST image is the `boot signature` of the DEST
+  `boot.img`.
+* The DEST retrofitted boot image must pass the `vts_gki_compliance_test`
+  testcase.
+
+### Retrofit to boot image V2
+
+* The `kernel` of the DEST image must be from the SOURCE `boot.img`.
+* The `ramdisk` of the DEST image must be from the SOURCE `vendor_boot.img` and
+  `init_boot.img`. The DEST `ramdisk` is the ramdisk concatenation of the vendor
+  ramdisk and generic ramdisk.
+* The `recovery dtbo / acpio` must be empty.
+* The `dtb` of the DEST image must be from the SOURCE `vendor_boot.img`.
+* The `boot_signature` section must be appended to the end of the boot image,
+  and its size is zero-padded to 16KiB.
+
+```
+  +---------------------+
+  | boot header         | 1 page
+  +---------------------+
+  | kernel              | n pages
+  +---------------------+
+  | * vendor ramdisk    |
+  |  +generic ramdisk   | m pages
+  +---------------------+
+  | second stage        | o pages
+  +---------------------+
+  | recovery dtbo/acpio | 0 byte
+  +---------------------+
+  | dtb                 | q pages
+  +---------------------+
+  | * boot signature    | 16384 (16K) bytes
+  +---------------------+
+```
diff --git a/gki/boot_signature_info.sh b/gki/boot_signature_info.sh
new file mode 100755
index 0000000..febeb1d
--- /dev/null
+++ b/gki/boot_signature_info.sh
@@ -0,0 +1,98 @@
+#!/bin/bash
+#
+# Copyright (C) 2022 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+#
+# Dump boot signature info of a GKI boot image.
+#
+
+set -eo errtrace
+
+die() {
+  echo >&2 "ERROR:" "${@}"
+  exit 1
+}
+
+TEMP_DIR="$(mktemp -d)"
+readonly TEMP_DIR
+
+exit_handler() {
+  readonly EXIT_CODE="$?"
+  rm -rf "${TEMP_DIR}" ||:
+  exit "${EXIT_CODE}"
+}
+
+trap exit_handler EXIT
+trap 'die "line ${LINENO}, ${FUNCNAME:-<main>}(): \"${BASH_COMMAND}\" returned \"$?\"" ' ERR
+
+get_arg() {
+  local arg="$1"
+  shift
+  while [[ "$#" -gt 0 ]]; do
+    if [[ "$1" == "${arg}" ]]; then
+      shift
+      echo "$1"
+      return
+    fi
+    shift
+  done
+}
+
+readonly VBMETA_IMAGE="${TEMP_DIR}/boot.boot_signature"
+readonly VBMETA_IMAGE_TEMP="${VBMETA_IMAGE}.temp"
+readonly VBMETA_INFO="${VBMETA_IMAGE}.info"
+readonly BOOT_IMAGE="${TEMP_DIR}/boot.img"
+readonly BOOT_IMAGE_DIR="${TEMP_DIR}/boot.unpack_dir"
+readonly BOOT_IMAGE_ARGS="${TEMP_DIR}/boot.mkbootimg_args"
+readonly BOOT_SIGNATURE_SIZE=$(( 16 << 10 ))
+
+[[ -f "$1" ]] ||
+  die "expected one input image"
+cp "$1" "${BOOT_IMAGE}"
+
+# This could fail if there already is no AVB footer.
+avbtool erase_footer --image "${BOOT_IMAGE}" 2>/dev/null ||:
+
+unpack_bootimg --boot_img "${BOOT_IMAGE}" --out "${BOOT_IMAGE_DIR}" \
+  --format=mkbootimg -0 > "${BOOT_IMAGE_ARGS}"
+
+declare -a boot_args=()
+while IFS= read -r -d '' ARG; do
+  boot_args+=("${ARG}")
+done < "${BOOT_IMAGE_ARGS}"
+
+BOOT_IMAGE_VERSION="$(get_arg --header_version "${boot_args[@]}")"
+if [[ "${BOOT_IMAGE_VERSION}" -ge 4 ]] && [[ -f "${BOOT_IMAGE_DIR}/boot_signature" ]]; then
+  cp "${BOOT_IMAGE_DIR}/boot_signature" "${VBMETA_IMAGE}"
+else
+  tail -c "${BOOT_SIGNATURE_SIZE}" "${BOOT_IMAGE}" > "${VBMETA_IMAGE}"
+fi
+
+# Keep carving out vbmeta image from the boot signature until we fail or EOF.
+# Failing is fine because there could be padding trailing the boot signature.
+while avbtool info_image --image "${VBMETA_IMAGE}" --output "${VBMETA_INFO}" 2>/dev/null; do
+  cat "${VBMETA_INFO}"
+  echo
+
+  declare -i H A X
+  H="$(cat "${VBMETA_INFO}" | grep 'Header Block:' | awk '{print $3}')"
+  A="$(cat "${VBMETA_INFO}" | grep 'Authentication Block:' | awk '{print $3}')"
+  X="$(cat "${VBMETA_INFO}" | grep 'Auxiliary Block:' | awk '{print $3}')"
+  vbmeta_size="$(( ${H} + ${A} + ${X} ))"
+
+  tail -c "+$(( ${vbmeta_size} + 1 ))" "${VBMETA_IMAGE}" > "${VBMETA_IMAGE_TEMP}"
+  cp "${VBMETA_IMAGE_TEMP}" "${VBMETA_IMAGE}"
+done
diff --git a/gki/certify_bootimg.py b/gki/certify_bootimg.py
new file mode 100755
index 0000000..9a7b058
--- /dev/null
+++ b/gki/certify_bootimg.py
@@ -0,0 +1,246 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022, The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Certify a GKI boot image by generating and appending its boot_signature."""
+
+from argparse import ArgumentParser
+import glob
+import os
+import shlex
+import shutil
+import subprocess
+import tempfile
+
+from gki.generate_gki_certificate import generate_gki_certificate
+from unpack_bootimg import unpack_bootimg
+
+BOOT_SIGNATURE_SIZE = 16 * 1024
+
+
+def get_kernel(boot_img):
+    """Extracts the kernel from |boot_img| and returns it."""
+    with tempfile.TemporaryDirectory() as unpack_dir:
+        unpack_bootimg(boot_img, unpack_dir)
+        with open(os.path.join(unpack_dir, 'kernel'), 'rb') as kernel:
+            kernel_bytes = kernel.read()
+            assert len(kernel_bytes) > 0
+            return kernel_bytes
+
+
+def add_certificate(boot_img, algorithm, key, extra_args):
+    """Appends certificates to the end of the boot image.
+
+    This functions appends two certificates to the end of the |boot_img|:
+    the 'boot' certificate and the 'generic_kernel' certificate. The former
+    is to certify the entire |boot_img|, while the latter is to certify
+    the kernel inside the |boot_img|.
+    """
+
+    def generate_certificate(image, certificate_name):
+        """Generates the certificate and returns the certificate content."""
+        with tempfile.NamedTemporaryFile() as output_certificate:
+            generate_gki_certificate(
+                image=image, avbtool='avbtool', name=certificate_name,
+                algorithm=algorithm, key=key, salt='d00df00d',
+                additional_avb_args=extra_args, output=output_certificate.name)
+            output_certificate.seek(os.SEEK_SET, 0)
+            return output_certificate.read()
+
+    boot_signature_bytes = b''
+    boot_signature_bytes += generate_certificate(boot_img, 'boot')
+
+    with tempfile.NamedTemporaryFile() as kernel_img:
+        kernel_img.write(get_kernel(boot_img))
+        kernel_img.flush()
+        boot_signature_bytes += generate_certificate(kernel_img.name,
+                                                     'generic_kernel')
+
+    if len(boot_signature_bytes) > BOOT_SIGNATURE_SIZE:
+        raise ValueError(
+            f'boot_signature size must be <= {BOOT_SIGNATURE_SIZE}')
+    boot_signature_bytes += (
+        b'\0' * (BOOT_SIGNATURE_SIZE - len(boot_signature_bytes)))
+    assert len(boot_signature_bytes) == BOOT_SIGNATURE_SIZE
+
+    with open(boot_img, 'ab') as f:
+        f.write(boot_signature_bytes)
+
+
+def erase_certificate_and_avb_footer(boot_img):
+    """Erases the boot certificate and avb footer.
+
+    A boot image might already contain a certificate and/or a AVB footer.
+    This function erases these additional metadata from the |boot_img|.
+    """
+    # Tries to erase the AVB footer first, which may or may not exist.
+    avbtool_cmd = ['avbtool', 'erase_footer', '--image', boot_img]
+    subprocess.run(avbtool_cmd, check=False, stderr=subprocess.DEVNULL)
+    assert os.path.getsize(boot_img) > 0
+
+    # No boot signature to erase, just return.
+    if os.path.getsize(boot_img) <= BOOT_SIGNATURE_SIZE:
+        return
+
+    # Checks if the last 16K is a boot signature, then erases it.
+    with open(boot_img, 'rb') as image:
+        image.seek(-BOOT_SIGNATURE_SIZE, os.SEEK_END)
+        boot_signature = image.read(BOOT_SIGNATURE_SIZE)
+        assert len(boot_signature) == BOOT_SIGNATURE_SIZE
+
+    with tempfile.NamedTemporaryFile() as signature_tmpfile:
+        signature_tmpfile.write(boot_signature)
+        signature_tmpfile.flush()
+        avbtool_info_cmd = [
+            'avbtool', 'info_image', '--image', signature_tmpfile.name]
+        result = subprocess.run(avbtool_info_cmd, check=False,
+                                stdout=subprocess.DEVNULL,
+                                stderr=subprocess.DEVNULL)
+        has_boot_signature = (result.returncode == 0)
+
+    if has_boot_signature:
+        new_file_size = os.path.getsize(boot_img) - BOOT_SIGNATURE_SIZE
+        os.truncate(boot_img, new_file_size)
+
+    assert os.path.getsize(boot_img) > 0
+
+
+def get_avb_image_size(image):
+    """Returns the image size if there is a AVB footer, else return zero."""
+
+    avbtool_info_cmd = ['avbtool', 'info_image', '--image', image]
+    result = subprocess.run(avbtool_info_cmd, check=False,
+                            stdout=subprocess.DEVNULL,
+                            stderr=subprocess.DEVNULL)
+
+    if result.returncode == 0:
+        return os.path.getsize(image)
+
+    return 0
+
+
+def add_avb_footer(image, partition_size):
+    """Appends a AVB hash footer to the image."""
+
+    avbtool_cmd = ['avbtool', 'add_hash_footer', '--image', image,
+                   '--partition_name', 'boot']
+
+    if partition_size:
+        avbtool_cmd.extend(['--partition_size', str(partition_size)])
+    else:
+        avbtool_cmd.extend(['--dynamic_partition_size'])
+
+    subprocess.check_call(avbtool_cmd)
+
+
+def load_dict_from_file(path):
+    """Loads key=value pairs from |path| and returns a dict."""
+    d = {}
+    with open(path, 'r', encoding='utf-8') as f:
+        for line in f:
+            line = line.strip()
+            if not line or line.startswith('#'):
+                continue
+            if '=' in line:
+                name, value = line.split('=', 1)
+                d[name] = value
+    return d
+
+
+def parse_cmdline():
+    """Parse command-line options."""
+    parser = ArgumentParser(add_help=True)
+
+    # Required args.
+    input_group = parser.add_mutually_exclusive_group(required=True)
+    input_group.add_argument(
+        '--boot_img', help='path to the boot image to certify')
+    input_group.add_argument(
+        '--boot_img_zip', help='path to the boot-img-*.zip archive to certify')
+
+    parser.add_argument('--algorithm', required=True,
+                        help='signing algorithm for the certificate')
+    parser.add_argument('--key', required=True,
+                        help='path to the RSA private key')
+    parser.add_argument('-o', '--output', required=True,
+                        help='output file name')
+
+    # Optional args.
+    parser.add_argument('--extra_args', default=[], action='append',
+                        help='extra arguments to be forwarded to avbtool')
+
+    args = parser.parse_args()
+
+    extra_args = []
+    for a in args.extra_args:
+        extra_args.extend(shlex.split(a))
+    args.extra_args = extra_args
+
+    return args
+
+
+def certify_bootimg(boot_img, output_img, algorithm, key, extra_args):
+    """Certify a GKI boot image by generating and appending a boot_signature."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        boot_tmp = os.path.join(temp_dir, 'boot.tmp')
+        shutil.copy2(boot_img, boot_tmp)
+
+        erase_certificate_and_avb_footer(boot_tmp)
+        add_certificate(boot_tmp, algorithm, key, extra_args)
+
+        avb_partition_size = get_avb_image_size(boot_img)
+        add_avb_footer(boot_tmp, avb_partition_size)
+
+        # We're done, copy the temp image to the final output.
+        shutil.copy2(boot_tmp, output_img)
+
+
+def certify_bootimg_zip(boot_img_zip, output_zip, algorithm, key, extra_args):
+    """Similar to certify_bootimg(), but for a zip archive of boot images."""
+    with tempfile.TemporaryDirectory() as unzip_dir:
+        shutil.unpack_archive(boot_img_zip, unzip_dir)
+
+        gki_info_file = os.path.join(unzip_dir, 'gki-info.txt')
+        if os.path.exists(gki_info_file):
+            info_dict = load_dict_from_file(gki_info_file)
+            if 'certify_bootimg_extra_args' in info_dict:
+                extra_args.extend(
+                    shlex.split(info_dict['certify_bootimg_extra_args']))
+
+        for boot_img in glob.glob(os.path.join(unzip_dir, 'boot-*.img')):
+            print(f'Certifying {os.path.basename(boot_img)} ...')
+            certify_bootimg(boot_img=boot_img, output_img=boot_img,
+                            algorithm=algorithm, key=key, extra_args=extra_args)
+
+        print(f'Making certified archive: {output_zip}')
+        archive_base_name = os.path.splitext(output_zip)[0]
+        shutil.make_archive(archive_base_name, 'zip', unzip_dir)
+
+
+def main():
+    """Parse arguments and certify the boot image."""
+    args = parse_cmdline()
+
+    if args.boot_img_zip:
+        certify_bootimg_zip(args.boot_img_zip, args.output, args.algorithm,
+                            args.key, args.extra_args)
+    else:
+        certify_bootimg(args.boot_img, args.output, args.algorithm,
+                        args.key, args.extra_args)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/gki/certify_bootimg_test.py b/gki/certify_bootimg_test.py
new file mode 100644
index 0000000..8c7c4d3
--- /dev/null
+++ b/gki/certify_bootimg_test.py
@@ -0,0 +1,773 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022, The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests certify_bootimg."""
+
+import logging
+import glob
+import os
+import random
+import shutil
+import struct
+import subprocess
+import sys
+import tempfile
+import unittest
+
+BOOT_SIGNATURE_SIZE = 16 * 1024
+
+TEST_KERNEL_CMDLINE = (
+    'printk.devkmsg=on firmware_class.path=/vendor/etc/ init=/init '
+    'kfence.sample_interval=500 loop.max_part=7 bootconfig'
+)
+
+
+def generate_test_file(pathname, size, seed=None):
+    """Generates a gibberish-filled test file and returns its pathname."""
+    random.seed(os.path.basename(pathname) if seed is None else seed)
+    with open(pathname, 'wb') as file:
+        file.write(random.randbytes(size))
+    return pathname
+
+
+def generate_test_boot_image(boot_img, kernel_size=4096, seed='kernel',
+                             avb_partition_size=None):
+    """Generates a test boot.img without a ramdisk."""
+    with tempfile.NamedTemporaryFile() as kernel_tmpfile:
+        generate_test_file(kernel_tmpfile.name, kernel_size, seed)
+        kernel_tmpfile.flush()
+
+        mkbootimg_cmds = [
+            'mkbootimg',
+            '--header_version', '4',
+            '--kernel', kernel_tmpfile.name,
+            '--cmdline', TEST_KERNEL_CMDLINE,
+            '--os_version', '12.0.0',
+            '--os_patch_level', '2022-03',
+            '--output', boot_img,
+        ]
+        subprocess.check_call(mkbootimg_cmds)
+
+    if avb_partition_size:
+        avbtool_cmd = ['avbtool', 'add_hash_footer', '--image', boot_img,
+                       '--partition_name', 'boot',
+                       '--partition_size', str(avb_partition_size)]
+        subprocess.check_call(avbtool_cmd)
+
+
+def generate_test_boot_image_archive(output_zip, boot_img_info, gki_info=None):
+    """Generates a zip archive of test boot images.
+
+    It also adds a file gki-info.txt, which contains additional settings for
+    for `certify_bootimg --extra_args`.
+
+    Args:
+        output_zip: the output zip archive, e.g., /path/to/boot-img.zip.
+        boot_img_info: a list of (boot_image_name, kernel_size,
+          partition_size) tuples. e.g.,
+          [('boot-1.0.img', 4096, 4 * 1024),
+           ('boot-2.0.img', 8192, 8 * 1024)].
+        gki_info: the file content to be written into 'gki-info.txt' in the
+          |output_zip|.
+    """
+    with tempfile.TemporaryDirectory() as temp_out_dir:
+        for name, kernel_size, partition_size in boot_img_info:
+            boot_img = os.path.join(temp_out_dir, name)
+            generate_test_boot_image(boot_img=boot_img,
+                                     kernel_size=kernel_size,
+                                     seed=name,
+                                     avb_partition_size=partition_size)
+
+        if gki_info:
+            gki_info_path = os.path.join(temp_out_dir, 'gki-info.txt')
+            with open(gki_info_path, 'w', encoding='utf-8') as f:
+                f.write(gki_info)
+
+        archive_base_name = os.path.splitext(output_zip)[0]
+        shutil.make_archive(archive_base_name, 'zip', temp_out_dir)
+
+
+def has_avb_footer(image):
+    """Returns true if the image has a avb footer."""
+
+    avbtool_info_cmd = ['avbtool', 'info_image', '--image', image]
+    result = subprocess.run(avbtool_info_cmd, check=False,
+                            stdout=subprocess.DEVNULL,
+                            stderr=subprocess.DEVNULL)
+
+    return result.returncode == 0
+
+
+def get_vbmeta_size(vbmeta_bytes):
+    """Returns the total size of a AvbVBMeta image."""
+
+    # Keep in sync with |AvbVBMetaImageHeader|.
+    AVB_MAGIC = b'AVB0'                        # pylint: disable=C0103
+    AVB_VBMETA_IMAGE_HEADER_SIZE = 256         # pylint: disable=C0103
+    FORMAT_STRING = (                          # pylint: disable=C0103
+        '!4s2L'      # magic, 2 x version.
+        '2Q'         # 2 x block size: Authentication and Auxiliary blocks.
+    )
+
+    if len(vbmeta_bytes) < struct.calcsize(FORMAT_STRING):
+        return 0
+
+    data = vbmeta_bytes[:struct.calcsize(FORMAT_STRING)]
+    (magic, _, _,
+     authentication_block_size,
+     auxiliary_data_block_size) = struct.unpack(FORMAT_STRING, data)
+
+    if magic == AVB_MAGIC:
+        return (AVB_VBMETA_IMAGE_HEADER_SIZE +
+                authentication_block_size +
+                auxiliary_data_block_size)
+    return 0
+
+
+def extract_boot_signatures(boot_img, output_dir):
+    """Extracts the boot signatures of a boot image.
+
+    This functions extracts the boot signatures of |boot_img| as:
+      - |output_dir|/boot_signature1
+      - |output_dir|/boot_signature2
+    """
+
+    boot_img_copy = os.path.join(output_dir, 'boot_image_copy')
+    shutil.copy2(boot_img, boot_img_copy)
+    avbtool_cmd = ['avbtool', 'erase_footer', '--image', boot_img_copy]
+    subprocess.run(avbtool_cmd, check=False, stderr=subprocess.DEVNULL)
+
+    # The boot signature is assumed to be at the end of boot image, after
+    # the AVB footer is erased.
+    with open(boot_img_copy, 'rb') as image:
+        image.seek(-BOOT_SIGNATURE_SIZE, os.SEEK_END)
+        boot_signature_bytes = image.read(BOOT_SIGNATURE_SIZE)
+        assert len(boot_signature_bytes) == BOOT_SIGNATURE_SIZE
+    os.unlink(boot_img_copy)
+
+    num_signatures = 0
+    while True:
+        next_signature_size = get_vbmeta_size(boot_signature_bytes)
+        if next_signature_size <= 0:
+            break
+
+        num_signatures += 1
+        next_signature = boot_signature_bytes[:next_signature_size]
+        output_path = os.path.join(
+            output_dir, 'boot_signature' + str(num_signatures))
+        with open(output_path, 'wb') as output:
+            output.write(next_signature)
+
+        # Moves to the next signature.
+        boot_signature_bytes = boot_signature_bytes[next_signature_size:]
+
+
+def extract_boot_archive_with_signatures(boot_img_zip, output_dir):
+    """Extracts boot images and signatures of a boot images archive.
+
+    Suppose there are two boot images in |boot_img_zip|: boot-1.0.img
+    and boot-2.0.img. This function then extracts each boot-*.img and
+    their signatures as:
+      - |output_dir|/boot-1.0.img
+      - |output_dir|/boot-2.0.img
+      - |output_dir|/boot-1.0/boot_signature1
+      - |output_dir|/boot-1.0/boot_signature2
+      - |output_dir|/boot-2.0/boot_signature1
+      - |output_dir|/boot-2.0/boot_signature2
+    """
+    shutil.unpack_archive(boot_img_zip, output_dir)
+    for boot_img in glob.glob(os.path.join(output_dir, 'boot-*.img')):
+        img_name = os.path.splitext(os.path.basename(boot_img))[0]
+        signature_output_dir = os.path.join(output_dir, img_name)
+        os.mkdir(signature_output_dir, 0o777)
+        extract_boot_signatures(boot_img, signature_output_dir)
+
+
+class CertifyBootimgTest(unittest.TestCase):
+    """Tests the functionalities of certify_bootimg."""
+
+    def setUp(self):
+        # Saves the test executable directory so that relative path references
+        # to test dependencies don't rely on being manually run from the
+        # executable directory.
+        # With this, we can just open "./testdata/testkey_rsa2048.pem" in the
+        # following tests with subprocess.run(..., cwd=self._exec_dir, ...).
+        self._exec_dir = os.path.abspath(os.path.dirname(sys.argv[0]))
+
+        # Set self.maxDiff to None to see full diff in assertion.
+        # C0103: invalid-name for maxDiff.
+        self.maxDiff = None  # pylint: disable=C0103
+
+        self._EXPECTED_BOOT_SIGNATURE_RSA2048 = (       # pylint: disable=C0103
+            'Minimum libavb version:   1.0\n'
+            'Header Block:             256 bytes\n'
+            'Authentication Block:     320 bytes\n'
+            'Auxiliary Block:          832 bytes\n'
+            'Public key (sha1):        '
+            'cdbb77177f731920bbe0a0f94f84d9038ae0617d\n'
+            'Algorithm:                SHA256_RSA2048\n'
+            'Rollback Index:           0\n'
+            'Flags:                    0\n'
+            'Rollback Index Location:  0\n'
+            "Release String:           'avbtool 1.2.0'\n"
+            'Descriptors:\n'
+            '    Hash descriptor:\n'
+            '      Image Size:            8192 bytes\n'
+            '      Hash Algorithm:        sha256\n'
+            '      Partition Name:        boot\n'           # boot
+            '      Salt:                  d00df00d\n'
+            '      Digest:                '
+            'faf1da72a4fba97ddab0b8f7a410db86'
+            '8fb72392a66d1440ff8bff490c73c771\n'
+            '      Flags:                 0\n'
+            "    Prop: gki -> 'nice'\n"
+            "    Prop: space -> 'nice to meet you'\n"
+        )
+
+        self._EXPECTED_KERNEL_SIGNATURE_RSA2048 = (     # pylint: disable=C0103
+            'Minimum libavb version:   1.0\n'
+            'Header Block:             256 bytes\n'
+            'Authentication Block:     320 bytes\n'
+            'Auxiliary Block:          832 bytes\n'
+            'Public key (sha1):        '
+            'cdbb77177f731920bbe0a0f94f84d9038ae0617d\n'
+            'Algorithm:                SHA256_RSA2048\n'
+            'Rollback Index:           0\n'
+            'Flags:                    0\n'
+            'Rollback Index Location:  0\n'
+            "Release String:           'avbtool 1.2.0'\n"
+            'Descriptors:\n'
+            '    Hash descriptor:\n'
+            '      Image Size:            4096 bytes\n'
+            '      Hash Algorithm:        sha256\n'
+            '      Partition Name:        generic_kernel\n' # generic_kernel
+            '      Salt:                  d00df00d\n'
+            '      Digest:                '
+            '762c877f3af0d50a4a4fbc1385d5c7ce'
+            '52a1288db74b33b72217d93db6f2909f\n'
+            '      Flags:                 0\n'
+            "    Prop: gki -> 'nice'\n"
+            "    Prop: space -> 'nice to meet you'\n"
+        )
+
+        self._EXPECTED_BOOT_SIGNATURE_RSA4096 = (       # pylint: disable=C0103
+            'Minimum libavb version:   1.0\n'
+            'Header Block:             256 bytes\n'
+            'Authentication Block:     576 bytes\n'
+            'Auxiliary Block:          1344 bytes\n'
+            'Public key (sha1):        '
+            '2597c218aae470a130f61162feaae70afd97f011\n'
+            'Algorithm:                SHA256_RSA4096\n'    # RSA4096
+            'Rollback Index:           0\n'
+            'Flags:                    0\n'
+            'Rollback Index Location:  0\n'
+            "Release String:           'avbtool 1.2.0'\n"
+            'Descriptors:\n'
+            '    Hash descriptor:\n'
+            '      Image Size:            8192 bytes\n'
+            '      Hash Algorithm:        sha256\n'
+            '      Partition Name:        boot\n'           # boot
+            '      Salt:                  d00df00d\n'
+            '      Digest:                '
+            'faf1da72a4fba97ddab0b8f7a410db86'
+            '8fb72392a66d1440ff8bff490c73c771\n'
+            '      Flags:                 0\n'
+            "    Prop: gki -> 'nice'\n"
+            "    Prop: space -> 'nice to meet you'\n"
+        )
+
+        self._EXPECTED_KERNEL_SIGNATURE_RSA4096 = (     # pylint: disable=C0103
+            'Minimum libavb version:   1.0\n'
+            'Header Block:             256 bytes\n'
+            'Authentication Block:     576 bytes\n'
+            'Auxiliary Block:          1344 bytes\n'
+            'Public key (sha1):        '
+            '2597c218aae470a130f61162feaae70afd97f011\n'
+            'Algorithm:                SHA256_RSA4096\n'    # RSA4096
+            'Rollback Index:           0\n'
+            'Flags:                    0\n'
+            'Rollback Index Location:  0\n'
+            "Release String:           'avbtool 1.2.0'\n"
+            'Descriptors:\n'
+            '    Hash descriptor:\n'
+            '      Image Size:            4096 bytes\n'
+            '      Hash Algorithm:        sha256\n'
+            '      Partition Name:        generic_kernel\n' # generic_kernel
+            '      Salt:                  d00df00d\n'
+            '      Digest:                '
+            '762c877f3af0d50a4a4fbc1385d5c7ce'
+            '52a1288db74b33b72217d93db6f2909f\n'
+            '      Flags:                 0\n'
+            "    Prop: gki -> 'nice'\n"
+            "    Prop: space -> 'nice to meet you'\n"
+        )
+
+        self._EXPECTED_BOOT_1_0_SIGNATURE1_RSA4096 = (   # pylint: disable=C0103
+            'Minimum libavb version:   1.0\n'
+            'Header Block:             256 bytes\n'
+            'Authentication Block:     576 bytes\n'
+            'Auxiliary Block:          1600 bytes\n'
+            'Public key (sha1):        '
+            '2597c218aae470a130f61162feaae70afd97f011\n'
+            'Algorithm:                SHA256_RSA4096\n'    # RSA4096
+            'Rollback Index:           0\n'
+            'Flags:                    0\n'
+            'Rollback Index Location:  0\n'
+            "Release String:           'avbtool 1.2.0'\n"
+            'Descriptors:\n'
+            '    Hash descriptor:\n'
+            '      Image Size:            12288 bytes\n'
+            '      Hash Algorithm:        sha256\n'
+            '      Partition Name:        boot\n'           # boot
+            '      Salt:                  d00df00d\n'
+            '      Digest:                '
+            '88465e463bffb9f7dfc0c1f46d01bcf3'
+            '15f7693e19bd188a0ca1feca2ed7b9df\n'
+            '      Flags:                 0\n'
+            "    Prop: gki -> 'nice'\n"
+            "    Prop: space -> 'nice to meet you'\n"
+            "    Prop: KERNEL_RELEASE -> '5.10.42-android13-0-00544-"
+            "ged21d463f856'\n"
+            "    Prop: BRANCH -> 'android13-5.10-2022-05'\n"
+            "    Prop: BUILD_NUMBER -> 'ab8295296'\n"
+            "    Prop: SPACE -> 'nice to meet you'\n"
+        )
+
+        self._EXPECTED_BOOT_1_0_SIGNATURE2_RSA4096 = (   # pylint: disable=C0103
+            'Minimum libavb version:   1.0\n'
+            'Header Block:             256 bytes\n'
+            'Authentication Block:     576 bytes\n'
+            'Auxiliary Block:          1600 bytes\n'
+            'Public key (sha1):        '
+            '2597c218aae470a130f61162feaae70afd97f011\n'
+            'Algorithm:                SHA256_RSA4096\n'    # RSA4096
+            'Rollback Index:           0\n'
+            'Flags:                    0\n'
+            'Rollback Index Location:  0\n'
+            "Release String:           'avbtool 1.2.0'\n"
+            'Descriptors:\n'
+            '    Hash descriptor:\n'
+            '      Image Size:            8192 bytes\n'
+            '      Hash Algorithm:        sha256\n'
+            '      Partition Name:        generic_kernel\n' # generic_kernel
+            '      Salt:                  d00df00d\n'
+            '      Digest:                '
+            '14ac8d0d233e57a317acd05cd458f2bb'
+            'cc78725ef9f66c1b38e90697fb09d943\n'
+            '      Flags:                 0\n'
+            "    Prop: gki -> 'nice'\n"
+            "    Prop: space -> 'nice to meet you'\n"
+            "    Prop: KERNEL_RELEASE -> '5.10.42-android13-0-00544-"
+            "ged21d463f856'\n"
+            "    Prop: BRANCH -> 'android13-5.10-2022-05'\n"
+            "    Prop: BUILD_NUMBER -> 'ab8295296'\n"
+            "    Prop: SPACE -> 'nice to meet you'\n"
+        )
+
+        self._EXPECTED_BOOT_2_0_SIGNATURE1_RSA4096 = (   # pylint: disable=C0103
+            'Minimum libavb version:   1.0\n'
+            'Header Block:             256 bytes\n'
+            'Authentication Block:     576 bytes\n'
+            'Auxiliary Block:          1600 bytes\n'
+            'Public key (sha1):        '
+            '2597c218aae470a130f61162feaae70afd97f011\n'
+            'Algorithm:                SHA256_RSA4096\n'    # RSA4096
+            'Rollback Index:           0\n'
+            'Flags:                    0\n'
+            'Rollback Index Location:  0\n'
+            "Release String:           'avbtool 1.2.0'\n"
+            'Descriptors:\n'
+            '    Hash descriptor:\n'
+            '      Image Size:            20480 bytes\n'
+            '      Hash Algorithm:        sha256\n'
+            '      Partition Name:        boot\n'           # boot
+            '      Salt:                  d00df00d\n'
+            '      Digest:                '
+            '3e6a9854a9d2350a7071083bc3f37376'
+            '37573fd87b1c72b146cb4870ac6af36f\n'
+            '      Flags:                 0\n'
+            "    Prop: gki -> 'nice'\n"
+            "    Prop: space -> 'nice to meet you'\n"
+            "    Prop: KERNEL_RELEASE -> '5.10.42-android13-0-00544-"
+            "ged21d463f856'\n"
+            "    Prop: BRANCH -> 'android13-5.10-2022-05'\n"
+            "    Prop: BUILD_NUMBER -> 'ab8295296'\n"
+            "    Prop: SPACE -> 'nice to meet you'\n"
+        )
+
+        self._EXPECTED_BOOT_2_0_SIGNATURE2_RSA4096 = (   # pylint: disable=C0103
+            'Minimum libavb version:   1.0\n'
+            'Header Block:             256 bytes\n'
+            'Authentication Block:     576 bytes\n'
+            'Auxiliary Block:          1600 bytes\n'
+            'Public key (sha1):        '
+            '2597c218aae470a130f61162feaae70afd97f011\n'
+            'Algorithm:                SHA256_RSA4096\n'    # RSA4096
+            'Rollback Index:           0\n'
+            'Flags:                    0\n'
+            'Rollback Index Location:  0\n'
+            "Release String:           'avbtool 1.2.0'\n"
+            'Descriptors:\n'
+            '    Hash descriptor:\n'
+            '      Image Size:            16384 bytes\n'
+            '      Hash Algorithm:        sha256\n'
+            '      Partition Name:        generic_kernel\n' # generic_kernel
+            '      Salt:                  d00df00d\n'
+            '      Digest:                '
+            '92fb8443cd284b67a4cbf5ce00348b50'
+            '1c657e0aedf4e2181c92ad7fc8b5224f\n'
+            '      Flags:                 0\n'
+            "    Prop: gki -> 'nice'\n"
+            "    Prop: space -> 'nice to meet you'\n"
+            "    Prop: KERNEL_RELEASE -> '5.10.42-android13-0-00544-"
+            "ged21d463f856'\n"
+            "    Prop: BRANCH -> 'android13-5.10-2022-05'\n"
+            "    Prop: BUILD_NUMBER -> 'ab8295296'\n"
+            "    Prop: SPACE -> 'nice to meet you'\n"
+        )
+
+        self._EXPECTED_BOOT_3_0_SIGNATURE1_RSA4096 = (   # pylint: disable=C0103
+            'Minimum libavb version:   1.0\n'
+            'Header Block:             256 bytes\n'
+            'Authentication Block:     576 bytes\n'
+            'Auxiliary Block:          1344 bytes\n'
+            'Public key (sha1):        '
+            '2597c218aae470a130f61162feaae70afd97f011\n'
+            'Algorithm:                SHA256_RSA4096\n'    # RSA4096
+            'Rollback Index:           0\n'
+            'Flags:                    0\n'
+            'Rollback Index Location:  0\n'
+            "Release String:           'avbtool 1.2.0'\n"
+            'Descriptors:\n'
+            '    Hash descriptor:\n'
+            '      Image Size:            12288 bytes\n'
+            '      Hash Algorithm:        sha256\n'
+            '      Partition Name:        boot\n'           # boot
+            '      Salt:                  d00df00d\n'
+            '      Digest:                '
+            '9b9cd845a367d7fc9b61d6ac02b0e7c9'
+            'dc3d3b219abf60dd6e19359f0353c917\n'
+            '      Flags:                 0\n'
+            "    Prop: gki -> 'nice'\n"
+            "    Prop: space -> 'nice to meet you'\n"
+        )
+
+        self._EXPECTED_BOOT_3_0_SIGNATURE2_RSA4096 = (   # pylint: disable=C0103
+            'Minimum libavb version:   1.0\n'
+            'Header Block:             256 bytes\n'
+            'Authentication Block:     576 bytes\n'
+            'Auxiliary Block:          1344 bytes\n'
+            'Public key (sha1):        '
+            '2597c218aae470a130f61162feaae70afd97f011\n'
+            'Algorithm:                SHA256_RSA4096\n'    # RSA4096
+            'Rollback Index:           0\n'
+            'Flags:                    0\n'
+            'Rollback Index Location:  0\n'
+            "Release String:           'avbtool 1.2.0'\n"
+            'Descriptors:\n'
+            '    Hash descriptor:\n'
+            '      Image Size:            8192 bytes\n'
+            '      Hash Algorithm:        sha256\n'
+            '      Partition Name:        generic_kernel\n' # generic_kernel
+            '      Salt:                  d00df00d\n'
+            '      Digest:                '
+            '0cd7d331ed9b32dcd92f00e2cac75595'
+            '52199170afe788a8fcf1954f9ea072d0\n'
+            '      Flags:                 0\n'
+            "    Prop: gki -> 'nice'\n"
+            "    Prop: space -> 'nice to meet you'\n"
+        )
+
+    def _test_boot_signatures(self, signatures_dir, expected_signatures_info):
+        """Tests the info of each boot signature under the signature directory.
+
+        Args:
+            signatures_dir: the directory containing the boot signatures. e.g.,
+                - signatures_dir/boot_signature1
+                - signatures_dir/boot_signature2
+            expected_signatures_info: A dict containing the expected output
+                of `avbtool info_image` for each signature under
+                |signatures_dir|. e.g.,
+                {'boot_signature1': expected_stdout_signature1
+                 'boot_signature2': expected_stdout_signature2}
+        """
+        for signature in expected_signatures_info:
+            avbtool_info_cmds = [
+                'avbtool', 'info_image', '--image',
+                os.path.join(signatures_dir, signature)
+            ]
+            result = subprocess.run(avbtool_info_cmds, check=True,
+                                    capture_output=True, encoding='utf-8')
+            self.assertEqual(result.stdout, expected_signatures_info[signature])
+
+    def test_certify_bootimg_without_avb_footer(self):
+        """Tests certify_bootimg on a boot image without an AVB footer."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            boot_img = os.path.join(temp_out_dir, 'boot.img')
+            generate_test_boot_image(boot_img)
+
+            # Generates the certified boot image, with a RSA2048 key.
+            boot_certified_img = os.path.join(temp_out_dir,
+                                              'boot-certified.img')
+            certify_bootimg_cmds = [
+                'certify_bootimg',
+                '--boot_img', boot_img,
+                '--algorithm', 'SHA256_RSA2048',
+                '--key', './testdata/testkey_rsa2048.pem',
+                '--extra_args', '--prop gki:nice '
+                '--prop space:"nice to meet you"',
+                '--output', boot_certified_img,
+            ]
+            subprocess.run(certify_bootimg_cmds, check=True, cwd=self._exec_dir)
+
+            extract_boot_signatures(boot_certified_img, temp_out_dir)
+            self._test_boot_signatures(
+                temp_out_dir,
+                {'boot_signature1': self._EXPECTED_BOOT_SIGNATURE_RSA2048,
+                 'boot_signature2': self._EXPECTED_KERNEL_SIGNATURE_RSA2048})
+
+            # Generates the certified boot image again, with a RSA4096 key.
+            boot_certified2_img = os.path.join(temp_out_dir,
+                                              'boot-certified2.img')
+            certify_bootimg_cmds = [
+                'certify_bootimg',
+                '--boot_img', boot_certified_img,
+                '--algorithm', 'SHA256_RSA4096',
+                '--key', './testdata/testkey_rsa4096.pem',
+                '--extra_args', '--prop gki:nice '
+                '--prop space:"nice to meet you"',
+                '--output', boot_certified2_img,
+            ]
+            subprocess.run(certify_bootimg_cmds, check=True, cwd=self._exec_dir)
+
+            extract_boot_signatures(boot_certified2_img, temp_out_dir)
+            self._test_boot_signatures(
+                temp_out_dir,
+                {'boot_signature1': self._EXPECTED_BOOT_SIGNATURE_RSA4096,
+                 'boot_signature2': self._EXPECTED_KERNEL_SIGNATURE_RSA4096})
+
+    def test_certify_bootimg_with_avb_footer(self):
+        """Tests the AVB footer location remains after certify_bootimg."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            boot_img = os.path.join(temp_out_dir, 'boot.img')
+            generate_test_boot_image(boot_img=boot_img,
+                                     avb_partition_size=128 * 1024)
+            self.assertTrue(has_avb_footer(boot_img))
+
+            # Generates the certified boot image, with a RSA2048 key.
+            boot_certified_img = os.path.join(temp_out_dir,
+                                              'boot-certified.img')
+            certify_bootimg_cmds = [
+                'certify_bootimg',
+                '--boot_img', boot_img,
+                '--algorithm', 'SHA256_RSA2048',
+                '--key', './testdata/testkey_rsa2048.pem',
+                '--extra_args', '--prop gki:nice '
+                '--prop space:"nice to meet you"',
+                '--output', boot_certified_img,
+            ]
+            subprocess.run(certify_bootimg_cmds, check=True, cwd=self._exec_dir)
+
+            # Checks an AVB footer exists and the image size remains.
+            self.assertTrue(has_avb_footer(boot_certified_img))
+            self.assertEqual(os.path.getsize(boot_img),
+                             os.path.getsize(boot_certified_img))
+
+            extract_boot_signatures(boot_certified_img, temp_out_dir)
+            self._test_boot_signatures(
+                temp_out_dir,
+                {'boot_signature1': self._EXPECTED_BOOT_SIGNATURE_RSA2048,
+                 'boot_signature2': self._EXPECTED_KERNEL_SIGNATURE_RSA2048})
+
+            # Generates the certified boot image again, with a RSA4096 key.
+            boot_certified2_img = os.path.join(temp_out_dir,
+                                              'boot-certified2.img')
+            certify_bootimg_cmds = [
+                'certify_bootimg',
+                '--boot_img', boot_certified_img,
+                '--algorithm', 'SHA256_RSA4096',
+                '--key', './testdata/testkey_rsa4096.pem',
+                '--extra_args', '--prop gki:nice '
+                '--prop space:"nice to meet you"',
+                '--output', boot_certified2_img,
+            ]
+            subprocess.run(certify_bootimg_cmds, check=True, cwd=self._exec_dir)
+
+            # Checks an AVB footer exists and the image size remains.
+            self.assertTrue(has_avb_footer(boot_certified2_img))
+            self.assertEqual(os.path.getsize(boot_certified_img),
+                             os.path.getsize(boot_certified2_img))
+
+            extract_boot_signatures(boot_certified2_img, temp_out_dir)
+            self._test_boot_signatures(
+                temp_out_dir,
+                {'boot_signature1': self._EXPECTED_BOOT_SIGNATURE_RSA4096,
+                 'boot_signature2': self._EXPECTED_KERNEL_SIGNATURE_RSA4096})
+
+    def test_certify_bootimg_exceed_size(self):
+        """Tests the boot signature size exceeded max size of the signature."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            boot_img = os.path.join(temp_out_dir, 'boot.img')
+            generate_test_boot_image(boot_img)
+
+            # Certifies the boot.img with many --extra_args, and checks
+            # it will raise the ValueError() exception.
+            boot_certified_img = os.path.join(temp_out_dir,
+                                              'boot-certified.img')
+            certify_bootimg_cmds = [
+                'certify_bootimg',
+                '--boot_img', boot_img,
+                '--algorithm', 'SHA256_RSA2048',
+                '--key', './testdata/testkey_rsa2048.pem',
+                # Makes it exceed the signature max size.
+                '--extra_args', '--prop foo:bar --prop gki:nice ' * 128,
+                '--output', boot_certified_img,
+            ]
+
+            try:
+                subprocess.run(certify_bootimg_cmds, check=True,
+                               capture_output=True, cwd=self._exec_dir,
+                               encoding='utf-8')
+                self.fail('Exceeding signature size assertion is not raised')
+            except subprocess.CalledProcessError as err:
+                self.assertIn('ValueError: boot_signature size must be <= ',
+                              err.stderr)
+
+    def test_certify_bootimg_archive(self):
+        """Tests certify_bootimg for a boot-img.zip."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            boot_img_zip = os.path.join(temp_out_dir, 'boot-img.zip')
+            gki_info = ('certify_bootimg_extra_args='
+                        '--prop KERNEL_RELEASE:5.10.42'
+                        '-android13-0-00544-ged21d463f856 '
+                        '--prop BRANCH:android13-5.10-2022-05 '
+                        '--prop BUILD_NUMBER:ab8295296 '
+                        '--prop SPACE:"nice to meet you"\n')
+            generate_test_boot_image_archive(
+                boot_img_zip,
+                # A list of (boot_img_name, kernel_size, partition_size).
+                [('boot-1.0.img', 8 * 1024, 128 * 1024),
+                 ('boot-2.0.img', 16 * 1024, 256 * 1024)],
+                gki_info)
+
+            # Certify the boot image archive, with a RSA4096 key.
+            boot_certified_img_zip = os.path.join(temp_out_dir,
+                                                  'boot-certified-img.zip')
+            certify_bootimg_cmds = [
+                'certify_bootimg',
+                '--boot_img_zip', boot_img_zip,
+                '--algorithm', 'SHA256_RSA4096',
+                '--key', './testdata/testkey_rsa4096.pem',
+                '--extra_args', '--prop gki:nice '
+                '--prop space:"nice to meet you"',
+                '--output', boot_certified_img_zip,
+            ]
+            subprocess.run(certify_bootimg_cmds, check=True, cwd=self._exec_dir)
+
+            extract_boot_archive_with_signatures(boot_certified_img_zip,
+                                                 temp_out_dir)
+
+            # Checks an AVB footer exists and the image size remains.
+            boot_1_img = os.path.join(temp_out_dir, 'boot-1.0.img')
+            self.assertTrue(has_avb_footer(boot_1_img))
+            self.assertEqual(os.path.getsize(boot_1_img), 128 * 1024)
+
+            boot_2_img = os.path.join(temp_out_dir, 'boot-2.0.img')
+            self.assertTrue(has_avb_footer(boot_2_img))
+            self.assertEqual(os.path.getsize(boot_2_img), 256 * 1024)
+
+            self._test_boot_signatures(
+                temp_out_dir,
+                {'boot-1.0/boot_signature1':
+                    self._EXPECTED_BOOT_1_0_SIGNATURE1_RSA4096,
+                 'boot-1.0/boot_signature2':
+                    self._EXPECTED_BOOT_1_0_SIGNATURE2_RSA4096,
+                 'boot-2.0/boot_signature1':
+                    self._EXPECTED_BOOT_2_0_SIGNATURE1_RSA4096,
+                 'boot-2.0/boot_signature2':
+                    self._EXPECTED_BOOT_2_0_SIGNATURE2_RSA4096})
+
+    def test_certify_bootimg_archive_without_gki_info(self):
+        """Tests certify_bootimg for a boot-img.zip."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            boot_img_zip = os.path.join(temp_out_dir, 'boot-img.zip')
+
+            # Checks ceritfy_bootimg works for a boot-img.zip without a
+            # gki-info.txt.
+            generate_test_boot_image_archive(
+                boot_img_zip,
+                # A list of (boot_img_name, kernel_size, partition_size).
+                [('boot-3.0.img', 8 * 1024, 128 * 1024)],
+                gki_info=None)
+            # Certify the boot image archive, with a RSA4096 key.
+            boot_certified_img_zip = os.path.join(temp_out_dir,
+                                                  'boot-certified-img.zip')
+            certify_bootimg_cmds = [
+                'certify_bootimg',
+                '--boot_img_zip', boot_img_zip,
+                '--algorithm', 'SHA256_RSA4096',
+                '--key', './testdata/testkey_rsa4096.pem',
+                '--extra_args', '--prop gki:nice '
+                '--prop space:"nice to meet you"',
+                '--output', boot_certified_img_zip,
+            ]
+            subprocess.run(certify_bootimg_cmds, check=True, cwd=self._exec_dir)
+
+            # Checks ceritfy_bootimg works for a boot-img.zip with a special
+            # gki-info.txt.
+            generate_test_boot_image_archive(
+                boot_img_zip,
+                # A list of (boot_img_name, kernel_size, partition_size).
+                [('boot-3.0.img', 8 * 1024, 128 * 1024)],
+                gki_info='a=b\n'
+                         'c=d\n')
+            # Certify the boot image archive, with a RSA4096 key.
+            boot_certified_img_zip = os.path.join(temp_out_dir,
+                                                  'boot-certified-img.zip')
+            certify_bootimg_cmds = [
+                'certify_bootimg',
+                '--boot_img_zip', boot_img_zip,
+                '--algorithm', 'SHA256_RSA4096',
+                '--key', './testdata/testkey_rsa4096.pem',
+                '--extra_args', '--prop gki:nice '
+                '--prop space:"nice to meet you"',
+                '--output', boot_certified_img_zip,
+            ]
+            subprocess.run(certify_bootimg_cmds, check=True, cwd=self._exec_dir)
+
+            extract_boot_archive_with_signatures(boot_certified_img_zip,
+                                                 temp_out_dir)
+
+            # Checks an AVB footer exists and the image size remains.
+            boot_3_img = os.path.join(temp_out_dir, 'boot-3.0.img')
+            self.assertTrue(has_avb_footer(boot_3_img))
+            self.assertEqual(os.path.getsize(boot_3_img), 128 * 1024)
+
+            self._test_boot_signatures(
+                temp_out_dir,
+                {'boot-3.0/boot_signature1':
+                    self._EXPECTED_BOOT_3_0_SIGNATURE1_RSA4096,
+                 'boot-3.0/boot_signature2':
+                    self._EXPECTED_BOOT_3_0_SIGNATURE2_RSA4096})
+
+
+# I don't know how, but we need both the logger configuration and verbosity
+# level > 2 to make atest work. And yes this line needs to be at the very top
+# level, not even in the "__main__" indentation block.
+logging.basicConfig(stream=sys.stdout)
+
+if __name__ == '__main__':
+    unittest.main(verbosity=2)
diff --git a/gki/generate_gki_certificate.py b/gki/generate_gki_certificate.py
new file mode 100755
index 0000000..2797cca
--- /dev/null
+++ b/gki/generate_gki_certificate.py
@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021, The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Generate a Generic Boot Image certificate suitable for VTS verification."""
+
+from argparse import ArgumentParser
+import shlex
+import subprocess
+
+
+def generate_gki_certificate(image, avbtool, name, algorithm, key, salt,
+                             additional_avb_args, output):
+    """Shell out to avbtool to generate a GKI certificate."""
+
+    # Need to specify a value of --partition_size for avbtool to work.
+    # We use 64 MB below, but avbtool will not resize the boot image to
+    # this size because --do_not_append_vbmeta_image is also specified.
+    avbtool_cmd = [
+        avbtool, 'add_hash_footer',
+        '--partition_name', name,
+        '--partition_size', str(64 * 1024 * 1024),
+        '--image', image,
+        '--algorithm', algorithm,
+        '--key', key,
+        '--do_not_append_vbmeta_image',
+        '--output_vbmeta_image', output,
+    ]
+
+    if salt is not None:
+        avbtool_cmd += ['--salt', salt]
+
+    avbtool_cmd += additional_avb_args
+
+    subprocess.check_call(avbtool_cmd)
+
+
+def parse_cmdline():
+    parser = ArgumentParser(add_help=True)
+
+    # Required args.
+    parser.add_argument('image', help='path to the image')
+    parser.add_argument('-o', '--output', required=True,
+                        help='output certificate file name')
+    parser.add_argument('--name', required=True,
+                        choices=['boot', 'generic_kernel'],
+                        help='name of the image to be certified')
+    parser.add_argument('--algorithm', required=True,
+                        help='AVB signing algorithm')
+    parser.add_argument('--key', required=True,
+                        help='path to the RSA private key')
+
+    # Optional args.
+    parser.add_argument('--avbtool', default='avbtool',
+                        help='path to the avbtool executable')
+    parser.add_argument('--salt', help='salt to use when computing image hash')
+    parser.add_argument('--additional_avb_args', default=[], action='append',
+                        help='additional arguments to be forwarded to avbtool')
+
+    args = parser.parse_args()
+
+    additional_avb_args = []
+    for a in args.additional_avb_args:
+        additional_avb_args.extend(shlex.split(a))
+    args.additional_avb_args = additional_avb_args
+
+    return args
+
+
+def main():
+    args = parse_cmdline()
+    generate_gki_certificate(
+        image=args.image, avbtool=args.avbtool, name=args.name,
+        algorithm=args.algorithm, key=args.key, salt=args.salt,
+        additional_avb_args=args.additional_avb_args,
+        output=args.output,
+    )
+
+
+if __name__ == '__main__':
+    main()
diff --git a/gki/retrofit_gki.sh b/gki/retrofit_gki.sh
new file mode 100755
index 0000000..01af7fa
--- /dev/null
+++ b/gki/retrofit_gki.sh
@@ -0,0 +1,231 @@
+#!/bin/bash
+#
+# Copyright (C) 2022 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+#
+# Retrofits GKI boot images for upgrading devices.
+#
+
+set -eo errtrace
+
+usage() {
+  cat <<EOF
+Usage:
+  $0 --boot BOOT --init_boot INIT_BOOT --version {3,4} -o OUTPUT
+  $0 --boot BOOT --init_boot INIT_BOOT --vendor_boot VENDOR_BOOT --version 2 -o OUTPUT
+
+Options:
+  --boot FILE
+    Path to the generic boot image.
+  --init_boot FILE
+    Path to the generic init_boot image.
+  --vendor_boot FILE
+    Path to the vendor boot image.
+  --version {2,3,4}
+    Boot image header version to retrofit to.
+  -o, --output FILE
+    Path to the output boot image.
+  -v, --verbose
+    Show debug messages.
+  -h, --help, --usage
+    Show this help message.
+EOF
+}
+
+die() {
+  echo >&2 "ERROR:" "${@}"
+  exit 1
+}
+
+file_size() {
+  stat -c '%s' "$1"
+}
+
+get_arg() {
+  local arg="$1"
+  shift
+  while [[ "$#" -gt 0 ]]; do
+    if [[ "$1" == "${arg}" ]]; then
+      shift
+      echo "$1"
+      return
+    fi
+    shift
+  done
+}
+
+TEMP_DIR="$(mktemp -d --tmpdir retrofit_gki.XXXXXXXX)"
+readonly TEMP_DIR
+
+exit_handler() {
+  readonly EXIT_CODE="$?"
+  rm -rf "${TEMP_DIR}" ||:
+  exit "${EXIT_CODE}"
+}
+
+trap exit_handler EXIT
+trap 'die "line ${LINENO}, ${FUNCNAME:-<main>}(): \"${BASH_COMMAND}\" returned \"$?\"" ' ERR
+
+while [[ "$1" =~ ^- ]]; do
+  case "$1" in
+    --boot )
+      shift
+      BOOT_IMAGE="$1"
+      ;;
+    --init_boot )
+      shift
+      INIT_BOOT_IMAGE="$1"
+      ;;
+    --vendor_boot )
+      shift
+      VENDOR_BOOT_IMAGE="$1"
+      ;;
+    --version )
+      shift
+      OUTPUT_BOOT_IMAGE_VERSION="$1"
+      ;;
+    -o | --output )
+      shift
+      OUTPUT_BOOT_IMAGE="$1"
+      ;;
+    -v | --verbose )
+      VERBOSE=true
+      ;;
+    -- )
+      shift
+      break
+      ;;
+    -h | --help | --usage )
+      usage
+      exit 0
+      ;;
+    * )
+      echo >&2 "Unexpected flag: '$1'"
+      usage >&2
+      exit 1
+      ;;
+  esac
+  shift
+done
+
+declare -ir OUTPUT_BOOT_IMAGE_VERSION
+readonly BOOT_IMAGE
+readonly INIT_BOOT_IMAGE
+readonly VENDOR_BOOT_IMAGE
+readonly OUTPUT_BOOT_IMAGE
+readonly VERBOSE
+
+# Make sure the input arguments make sense.
+[[ -f "${BOOT_IMAGE}" ]] ||
+  die "argument '--boot': not a regular file: '${BOOT_IMAGE}'"
+[[ -f "${INIT_BOOT_IMAGE}" ]] ||
+  die "argument '--init_boot': not a regular file: '${INIT_BOOT_IMAGE}'"
+if [[ "${OUTPUT_BOOT_IMAGE_VERSION}" -lt 2 ]] || [[ "${OUTPUT_BOOT_IMAGE_VERSION}" -gt 4 ]]; then
+  die "argument '--version': valid choices are {2, 3, 4}"
+elif [[ "${OUTPUT_BOOT_IMAGE_VERSION}" -eq 2 ]]; then
+  [[ -f "${VENDOR_BOOT_IMAGE}" ]] ||
+    die "argument '--vendor_boot': not a regular file: '${VENDOR_BOOT_IMAGE}'"
+fi
+[[ -z "${OUTPUT_BOOT_IMAGE}" ]] &&
+  die "argument '--output': cannot be empty"
+
+readonly BOOT_IMAGE_WITHOUT_AVB_FOOTER="${TEMP_DIR}/boot.img.without_avb_footer"
+readonly BOOT_DIR="${TEMP_DIR}/boot"
+readonly INIT_BOOT_DIR="${TEMP_DIR}/init_boot"
+readonly VENDOR_BOOT_DIR="${TEMP_DIR}/vendor_boot"
+readonly VENDOR_BOOT_MKBOOTIMG_ARGS="${TEMP_DIR}/vendor_boot.mkbootimg_args"
+readonly OUTPUT_RAMDISK="${TEMP_DIR}/out.ramdisk"
+readonly OUTPUT_BOOT_SIGNATURE="${TEMP_DIR}/out.boot_signature"
+
+readonly AVBTOOL="${AVBTOOL:-avbtool}"
+readonly MKBOOTIMG="${MKBOOTIMG:-mkbootimg}"
+readonly UNPACK_BOOTIMG="${UNPACK_BOOTIMG:-unpack_bootimg}"
+
+# Fixed boot signature size for easy discovery in VTS.
+readonly BOOT_SIGNATURE_SIZE=$(( 16 << 10 ))
+
+
+#
+# Preparations are done. Now begin the actual work.
+#
+
+# Copy the boot image because `avbtool erase_footer` edits the file in-place.
+cp "${BOOT_IMAGE}" "${BOOT_IMAGE_WITHOUT_AVB_FOOTER}"
+( [[ -n "${VERBOSE}" ]] && set -x
+  "${AVBTOOL}" erase_footer --image "${BOOT_IMAGE_WITHOUT_AVB_FOOTER}" 2>/dev/null ||:
+  tail -c "${BOOT_SIGNATURE_SIZE}" "${BOOT_IMAGE_WITHOUT_AVB_FOOTER}" > "${OUTPUT_BOOT_SIGNATURE}"
+  "${UNPACK_BOOTIMG}" --boot_img "${BOOT_IMAGE}" --out "${BOOT_DIR}" >/dev/null
+  "${UNPACK_BOOTIMG}" --boot_img "${INIT_BOOT_IMAGE}" --out "${INIT_BOOT_DIR}" >/dev/null
+)
+if [[ "$(file_size "${OUTPUT_BOOT_SIGNATURE}")" -ne "${BOOT_SIGNATURE_SIZE}" ]]; then
+  die "boot signature size must be equal to ${BOOT_SIGNATURE_SIZE}"
+fi
+
+declare -a mkbootimg_args=()
+
+if [[ "${OUTPUT_BOOT_IMAGE_VERSION}" -eq 4 ]]; then
+  mkbootimg_args+=( \
+    --header_version 4 \
+    --kernel "${BOOT_DIR}/kernel" \
+    --ramdisk "${INIT_BOOT_DIR}/ramdisk" \
+  )
+elif [[ "${OUTPUT_BOOT_IMAGE_VERSION}" -eq 3 ]]; then
+  mkbootimg_args+=( \
+    --header_version 3 \
+    --kernel "${BOOT_DIR}/kernel" \
+    --ramdisk "${INIT_BOOT_DIR}/ramdisk" \
+  )
+elif [[ "${OUTPUT_BOOT_IMAGE_VERSION}" -eq 2 ]]; then
+  ( [[ -n "${VERBOSE}" ]] && set -x
+    "${UNPACK_BOOTIMG}" --boot_img "${VENDOR_BOOT_IMAGE}" --out "${VENDOR_BOOT_DIR}" \
+      --format=mkbootimg -0 > "${VENDOR_BOOT_MKBOOTIMG_ARGS}"
+    cat "${VENDOR_BOOT_DIR}/vendor_ramdisk" "${INIT_BOOT_DIR}/ramdisk" > "${OUTPUT_RAMDISK}"
+  )
+
+  declare -a vendor_boot_args=()
+  while IFS= read -r -d '' ARG; do
+    vendor_boot_args+=("${ARG}")
+  done < "${VENDOR_BOOT_MKBOOTIMG_ARGS}"
+
+  pagesize="$(get_arg --pagesize "${vendor_boot_args[@]}")"
+  kernel_offset="$(get_arg --kernel_offset "${vendor_boot_args[@]}")"
+  ramdisk_offset="$(get_arg --ramdisk_offset "${vendor_boot_args[@]}")"
+  tags_offset="$(get_arg --tags_offset "${vendor_boot_args[@]}")"
+  dtb_offset="$(get_arg --dtb_offset "${vendor_boot_args[@]}")"
+  kernel_cmdline="$(get_arg --vendor_cmdline "${vendor_boot_args[@]}")"
+
+  mkbootimg_args+=( \
+    --header_version 2 \
+    --base 0 \
+    --kernel_offset "${kernel_offset}" \
+    --ramdisk_offset "${ramdisk_offset}" \
+    --second_offset 0 \
+    --tags_offset "${tags_offset}" \
+    --dtb_offset "${dtb_offset}" \
+    --cmdline "${kernel_cmdline}" \
+    --pagesize "${pagesize}" \
+    --kernel "${BOOT_DIR}/kernel" \
+    --ramdisk "${OUTPUT_RAMDISK}" \
+  )
+  if [[ -f "${VENDOR_BOOT_DIR}/dtb" ]]; then
+    mkbootimg_args+=(--dtb "${VENDOR_BOOT_DIR}/dtb")
+  fi
+fi
+
+( [[ -n "${VERBOSE}" ]] && set -x
+  "${MKBOOTIMG}" "${mkbootimg_args[@]}" --output "${OUTPUT_BOOT_IMAGE}"
+  cat "${OUTPUT_BOOT_SIGNATURE}" >> "${OUTPUT_BOOT_IMAGE}"
+)
diff --git a/gki/retrofit_gki_test.sh b/gki/retrofit_gki_test.sh
new file mode 100755
index 0000000..b3cb0a5
--- /dev/null
+++ b/gki/retrofit_gki_test.sh
@@ -0,0 +1,144 @@
+#!/bin/bash
+#
+# Copyright (C) 2022 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+set -eo errtrace
+
+die() {
+  echo >&2 "ERROR:" "${@}"
+  exit 1
+}
+
+trap 'die "line ${LINENO}, ${FUNCNAME:-<main>}(): \"${BASH_COMMAND}\" returned \"$?\"" ' ERR
+
+# Figure out where we are and where to look for test executables.
+cd "$(dirname "${BASH_SOURCE[0]}")"
+TEST_DIR="$(pwd)"
+readonly TEST_DIR
+readonly TEMP_DIR="${TEST_DIR}/stage.retrofit_gki_test"
+
+export PATH="${TEST_DIR}:${PATH}"
+rm -rf "${TEMP_DIR}"
+mkdir -p "${TEMP_DIR}"
+
+# Generate some test files.
+readonly TEST_DTB="${TEMP_DIR}/dtb"
+readonly TEST_KERNEL="${TEMP_DIR}/kernel"
+readonly TEST_RAMDISK="${TEMP_DIR}/ramdisk"
+readonly TEST_VENDOR_RAMDISK="${TEMP_DIR}/vendor_ramdisk"
+readonly TEST_BOOT_SIGNATURE="${TEMP_DIR}/boot.boot_signature"
+readonly TEST_V2_RETROFITTED_RAMDISK="${TEMP_DIR}/retrofitted.ramdisk"
+readonly TEST_BOOT_IMAGE="${TEMP_DIR}/boot.img"
+readonly TEST_INIT_BOOT_IMAGE="${TEMP_DIR}/init_boot.img"
+readonly TEST_VENDOR_BOOT_IMAGE="${TEMP_DIR}/vendor_boot.img"
+
+( # Run these in subshell because dd is noisy.
+  dd if=/dev/urandom of="${TEST_DTB}" bs=1024 count=10
+  dd if=/dev/urandom of="${TEST_KERNEL}" bs=1024 count=10
+  dd if=/dev/urandom of="${TEST_RAMDISK}" bs=1024 count=10
+  dd if=/dev/urandom of="${TEST_VENDOR_RAMDISK}" bs=1024 count=10
+  dd if=/dev/urandom of="${TEST_BOOT_SIGNATURE}" bs=1024 count=16
+) 2> /dev/null
+
+cat "${TEST_VENDOR_RAMDISK}" "${TEST_RAMDISK}" > "${TEST_V2_RETROFITTED_RAMDISK}"
+
+mkbootimg \
+  --header_version 4 \
+  --kernel "${TEST_KERNEL}" \
+  --output "${TEST_BOOT_IMAGE}"
+cat "${TEST_BOOT_SIGNATURE}" >> "${TEST_BOOT_IMAGE}"
+avbtool add_hash_footer --image "${TEST_BOOT_IMAGE}" --partition_name boot --partition_size $((20 << 20))
+
+mkbootimg \
+  --header_version 4 \
+  --ramdisk "${TEST_RAMDISK}" \
+  --output "${TEST_INIT_BOOT_IMAGE}"
+mkbootimg \
+  --header_version 4 \
+  --pagesize 4096 \
+  --dtb "${TEST_DTB}" \
+  --vendor_ramdisk "${TEST_VENDOR_RAMDISK}" \
+  --vendor_boot "${TEST_VENDOR_BOOT_IMAGE}"
+
+readonly RETROFITTED_IMAGE="${TEMP_DIR}/retrofitted_boot.img"
+readonly RETROFITTED_IMAGE_DIR="${TEMP_DIR}/retrofitted_boot.img.unpack"
+readonly BOOT_SIGNATURE_SIZE=$(( 16 << 10 ))
+
+
+#
+# Begin test.
+#
+echo >&2 "TEST: retrofit to boot v4"
+
+retrofit_gki.sh \
+  --boot "${TEST_BOOT_IMAGE}" \
+  --init_boot "${TEST_INIT_BOOT_IMAGE}" \
+  --version 4 \
+  --output "${RETROFITTED_IMAGE}"
+
+rm -rf "${RETROFITTED_IMAGE_DIR}"
+unpack_bootimg --boot_img "${RETROFITTED_IMAGE}" --out "${RETROFITTED_IMAGE_DIR}" > /dev/null
+tail -c "${BOOT_SIGNATURE_SIZE}" "${RETROFITTED_IMAGE}" > "${RETROFITTED_IMAGE_DIR}/boot_signature"
+
+cmp -s "${TEST_KERNEL}" "${RETROFITTED_IMAGE_DIR}/kernel" ||
+  die "unexpected diff: kernel"
+cmp -s "${TEST_RAMDISK}" "${RETROFITTED_IMAGE_DIR}/ramdisk" ||
+  die "unexpected diff: ramdisk"
+cmp -s "${TEST_BOOT_SIGNATURE}" "${RETROFITTED_IMAGE_DIR}/boot_signature" ||
+  die "unexpected diff: boot signature"
+
+
+echo >&2 "TEST: retrofit to boot v3"
+
+retrofit_gki.sh \
+  --boot "${TEST_BOOT_IMAGE}" \
+  --init_boot "${TEST_INIT_BOOT_IMAGE}" \
+  --version 3 \
+  --output "${RETROFITTED_IMAGE}"
+
+rm -rf "${RETROFITTED_IMAGE_DIR}"
+unpack_bootimg --boot_img "${RETROFITTED_IMAGE}" --out "${RETROFITTED_IMAGE_DIR}" > /dev/null
+tail -c "${BOOT_SIGNATURE_SIZE}" "${RETROFITTED_IMAGE}" > "${RETROFITTED_IMAGE_DIR}/boot_signature"
+
+cmp -s "${TEST_KERNEL}" "${RETROFITTED_IMAGE_DIR}/kernel" ||
+  die "unexpected diff: kernel"
+cmp -s "${TEST_RAMDISK}" "${RETROFITTED_IMAGE_DIR}/ramdisk" ||
+  die "unexpected diff: ramdisk"
+cmp -s "${TEST_BOOT_SIGNATURE}" "${RETROFITTED_IMAGE_DIR}/boot_signature" ||
+  die "unexpected diff: boot signature"
+
+
+echo >&2 "TEST: retrofit to boot v2"
+
+retrofit_gki.sh \
+  --boot "${TEST_BOOT_IMAGE}" \
+  --init_boot "${TEST_INIT_BOOT_IMAGE}" \
+  --vendor_boot "${TEST_VENDOR_BOOT_IMAGE}" \
+  --version 2 \
+  --output "${RETROFITTED_IMAGE}"
+
+rm -rf "${RETROFITTED_IMAGE_DIR}"
+unpack_bootimg --boot_img "${RETROFITTED_IMAGE}" --out "${RETROFITTED_IMAGE_DIR}" > /dev/null
+tail -c "${BOOT_SIGNATURE_SIZE}" "${RETROFITTED_IMAGE}" > "${RETROFITTED_IMAGE_DIR}/boot_signature"
+
+cmp -s "${TEST_DTB}" "${RETROFITTED_IMAGE_DIR}/dtb" ||
+  die "unexpected diff: dtb"
+cmp -s "${TEST_KERNEL}" "${RETROFITTED_IMAGE_DIR}/kernel" ||
+  die "unexpected diff: kernel"
+cmp -s "${TEST_V2_RETROFITTED_RAMDISK}" "${RETROFITTED_IMAGE_DIR}/ramdisk" ||
+  die "unexpected diff: ramdisk"
+cmp -s "${TEST_BOOT_SIGNATURE}" "${RETROFITTED_IMAGE_DIR}/boot_signature" ||
+  die "unexpected diff: boot signature"
diff --git a/gki/testdata/testkey_rsa2048.pem b/gki/testdata/testkey_rsa2048.pem
new file mode 100644
index 0000000..867dcff
--- /dev/null
+++ b/gki/testdata/testkey_rsa2048.pem
@@ -0,0 +1,27 @@
+-----BEGIN RSA PRIVATE KEY-----
+MIIEowIBAAKCAQEAxlVR3TIkouAOvH79vaJTgFhpfvVKQIeVkFRZPVXK/zY0Gvrh
+4JAqGjJoW/PfrQv5sdD36qtHH3a+G5hLZ6Ni+t/mtfjucxZfuLGC3kmJ1T3XqEKZ
+gXXI2IR7vVSoImREvDQGEDyJwtHzLANlkbGg0cghVhWZSCAndO8BenalC2v94/rt
+DfkPekH6dgU3Sf40T0sBSeSY94mOzTaqOR2pfV1rWlLRdWmo33zeHBv52Rlbt0dM
+uXAureXWiHztkm5GCBC1dgM+CaxNtizNEgC91KcD0xuRCCM2WxH+r1lpszyIJDct
+YbrFmVEYl/kjQpafhy7Nsk1fqSTyRdriZSYmTQIDAQABAoIBAQC+kJgaCuX8wYAn
+SXWQ0fmdZlXnMNRpcF0a0pD0SAzGb1RdYBXMaXiqtyhiwc53PPxsCDdNecjayIMd
+jJVXPTwLhTruOgMS/bp3gcgWwV34UHV4LJXGOGAE+jbS0hbDBMiudOYmj6RmVshp
+z9G1zZCSQNMXHaWsEYkX59XpzzoB384nRul2QgEtwzUNR9XlpzgtJBLk3SACkvsN
+mQ/DW8IWHXLg8vLn1LzVJ2e3B16H4MoE2TCHxqfMgr03IDRRJogkenQuQsFhevYT
+o/mJyHSWavVgzMHG9I5m+eepF4Wyhj1Y4WyKAuMI+9dHAX/h7Lt8XFCQCh5DbkVG
+zGr34sWBAoGBAOs7n7YZqNaaguovfIdRRsxxZr1yJAyDsr6w3yGImDZYju4c4WY9
+5esO2kP3FA4p0c7FhQF5oOb1rBuHEPp36cpL4aGeK87caqTfq63WZAujoTZpr9Lp
+BRbkL7w/xG7jpQ/clpA8sHzHGQs/nelxoOtC7E118FiRgvD/jdhlMyL9AoGBANfX
+vyoN1pplfT2xR8QOjSZ+Q35S/+SAtMuBnHx3l0qH2bbBjcvM1MNDWjnRDyaYhiRu
+i+KA7tqfib09+XpB3g5D6Ov7ls/Ldx0S/VcmVWtia2HK8y8iLGtokoBZKQ5AaFX2
+iQU8+tC4h69GnJYQKqNwgCUzh8+gHX5Y46oDiTmRAoGAYpOx8lX+czB8/Da6MNrW
+mIZNT8atZLEsDs2ANEVRxDSIcTCZJId7+m1W+nRoaycLTWNowZ1+2ErLvR10+AGY
+b7Ys79Wg9idYaY9yGn9lnZsMzAiuLeyIvXcSqgjvAKlVWrhOQFOughvNWvFl85Yy
+oWSCMlPiTLtt7CCsCKsgKuECgYBgdIp6GZsIfkgclKe0hqgvRoeU4TR3gcjJlM9A
+lBTo+pKhaBectplx9RxR8AnsPobbqwcaHnIfAuKDzjk5mEvKZjClnFXF4HAHbyAF
+nRzZEy9XkWFhc80T5rRpZO7C7qdxmu2aiKixM3V3L3/0U58qULEDbubHMw9bEhAT
+PudI8QKBgHEEiMm/hr9T41hbQi/LYanWnlFw1ue+osKuF8bXQuxnnHNuFT/c+9/A
+vWhgqG6bOEHu+p/IPrYm4tBMYlwsyh4nXCyGgDJLbLIfzKwKAWCtH9LwnyDVhOow
+GH9shdR+sW3Ew97xef02KAH4VlNANEmBV4sQNqWWvsYrcFm2rOdL
+-----END RSA PRIVATE KEY-----
diff --git a/gki/testdata/testkey_rsa4096.pem b/gki/testdata/testkey_rsa4096.pem
new file mode 100644
index 0000000..26db5c3
--- /dev/null
+++ b/gki/testdata/testkey_rsa4096.pem
@@ -0,0 +1,51 @@
+-----BEGIN RSA PRIVATE KEY-----
+MIIJKQIBAAKCAgEA2ASv49OEbH4NiT3CjNMSVeliyfEPXswWcqtEfCxlSpS1FisA
+uwbvEwdTTPlkuSh6G4SYiNhnpCP5p0vcSg/3OhiuVKgV/rCtrDXaO60nvK/o0y83
+NNZRK2xaJ9eWBq9ruIDK+jC0sYWzTaqqwxY0Grjnx/r5CXerl5PrRK7PILzwgBHb
+IwxHcblt1ntgR4cWVpO3wiqasEwBDDDYk4fw7W6LvjBb9qav3YB8RV6PkZNeRP64
+ggfuecq/MXNiWOPNxLzCER2hSr/+J32h9jWjXsrcVy8+8Mldhmr4r2an7c247aFf
+upuFGtUJrpROO8/LXMl5gPfMpkqoatjTMRH59gJjKhot0RpmGxZBvb33TcBK5SdJ
+X39Y4yct5clmDlI4Fjj7FutTP+b96aJeJVnYeUX/A0wmogBajsJRoRX5e/RcgZsY
+RzXYLQXprQ81dBWjjovMJ9p8XeT6BNMFC7o6sklFL0fHDUE/l4BNP8G1u3Bfpzev
+SCISRS71D4eS4oQB+RIPFBUkzomZ7rnEF3BwFeq+xmwfYrP0LRaH+1YeRauuMuRe
+ke1TZl697a3mEjkNg8noa2wtpe7EWmaujJfXDWxJx/XEkjGLCe4z2qk3tkkY+A5g
+Rcgzke8gVxC+eC2DJtbKYfkv4L8FMFJaEhwAp13MfC7FlYujO/BDLl7dANsCAwEA
+AQKCAgAWoL8P/WsktjuSwb5sY/vKtgzcHH1Ar942GsysuTXPDy686LpF3R8T/jNy
+n7k2UBAia8xSoWCR6BbRuHeV5oA+PLGeOpE7QaSfonB+yc+cy0x3Or3ssfqEsu/q
+toGHp75/8DXS6WE0K04x94u1rdC9b9sPrrGBlWCLGzqM0kbuJfyHXdd3n2SofAUO
+b5QRSgxD+2tHUpEroHqHnWJCaf4J0QegX45yktlfOYNK/PHLDQXV8ly/ejc32M4Y
+Tv7hUtOOJTuq8VCg9OWZm2Zo1QuM9XEJTPCp5l3+o5vzO6yhk2gotDvD32CdA+3k
+tLJRP54M1Sn+IXb1gGKN9rKAtGJbenWIPlNObhQgkbwG89Qd+5rfMXsiPv1Hl1tK
++tqwjD82/H3/ElaaMnwHCpeoGSp95OblAoBjzjMP2KsbvKSdL8O/rf1c3uOw9+DF
+cth0SA8y3ZzI11gJtb2QMGUrCny5n4sPGGbc3x38NdLhwbkPKZy60OiT4g2kNpdY
+dIitmAML2otttiF4AJM6AraPk8YVzkPLTksoL3azPBya5lIoDI2H3QvTtSvpXkXP
+yKchsDSWYbdqfplqC/X0Djp2/Zd8jpN5I6+1aSmpTmbwx/JTllY1N89FRZLIdxoh
+2k81LPiXhE6uRbjioJUlbnEWIpY2y2N2Clmxpjh0/IcXd1XImQKCAQEA7Zai+yjj
+8xit24aO9Tf3mZBXBjSaDodjC2KS1yCcAIXp6S7aH0wZipyZpQjys3zaBQyMRYFG
+bQqIfVAa6inWyDoofbAJHMu5BVcHFBPZvSS5YhDjc8XZ5dqSCxzIz9opIqAbm+b4
+aEV/3A3Jki5Dy8y/5j21GAK4Y4mqQOYzne7bDGi3Hyu041MGM4qfIcIkS5N1eHW4
+sDZJh6+K5tuxN5TX3nDZSpm9luNH8mLGgKAZ15b1LqXAtM5ycoBY9Hv082suPPom
+O+r0ybdRX6nDSH8+11y2KiP2kdVIUHCGkwlqgrux5YZyjCZPwOvEPhzSoOS+vBiF
+UVXA8idnxNLk1QKCAQEA6MIihDSXx+350fWqhQ/3Qc6gA/t2C15JwJ9+uFWA+gjd
+c/hn5HcmnmBJN4R04nLG/aU9SQur87a4mnC/Mp9JIARjHlZ/WNT4U0sJyPEVRg5U
+Z9VajAucWwi0JyJYCO1EMMy68Jp8qlTriK/L7nbD86JJ5ASxjojiN/0psK/Pk60F
+Rr+shKPi3jRQ1BDjDtAxOfo4ctf/nFbUM4bY0FNPQMP7WesoSKU0NBCRR6d0d2tq
+YflMjIQHx+N74P5jEdSCHTVGQm+dj47pUt3lLPLWc0bX1G/GekwXP4NUsR/70Hsi
+bwxkNnK2TSGzkt2rcOnutP125rJu6WpV7SNrq9rm7wKCAQAfMROcnbWviKHqnDPQ
+hdR/2K9UJTvEhInASOS2UZWpi+s1rez9BuSjigOx4wbaAZ4t44PW7C3uyt84dHfU
+HkIQb3I5bg8ENMrJpK9NN33ykwuzkDwMSwFcZ+Gci97hSubzoMl/IkeiiN1MapL4
+GhLUgsD+3UMVL+Y9SymK8637IgyoCGdiND6/SXsa8SwLJo3VTjqx4eKpX7cvlSBL
+RrRxc50TmwUsAhsd4CDl9YnSATLjVvJBeYlfM2tbFPaYwl1aR8v+PWkfnK0efm60
+fHki33HEnGteBPKuGq4vwVYpn6bYGwQz+f6335/A2DMfZHFSpjVURHPcRcHbCMla
+0cUxAoIBAQC25eYNkO478mo+bBbEXJlkoqLmvjAyGrNFo48F9lpVH6Y0vNuWkXJN
+PUgLUhAu6RYotjGENqG17rz8zt/PPY9Ok2P3sOx8t00y1mIn/hlDZXs55FM0fOMu
+PZaiscAPs7HDzvyOmDah+fzi+ZD8H2M3DS2W+YE0iaeJa2vZJS2t02W0BGXiDI33
+IZDqMyLYvwwPjOnShJydEzXID4xLl0tNjzLxo3GSNA7jYqlmbtV8CXIc7rMSL6WV
+ktIDKKJcnmpn3TcKeX6MEjaSIT82pNOS3fY3PmXuL+CMzfw8+u77Eecq78fHaTiL
+P5JGM93F6mzi19EY0tmInUBMCWtQLcENAoIBAQCg0KaOkb8T36qzPrtgbfou0E2D
+ufdpL1ugmD4edOFKQB5fDFQhLnSEVSJq3KUg4kWsXapQdsBd6kLdxS+K6MQrLBzr
+4tf0c7UCF1AzWk6wXMExZ8mRb2RkGZYQB2DdyhFB3TPmnq9CW8JCq+6kxg/wkU4s
+vM4JXzgcqVoSf42QJl+B9waeWhg0BTWx01lal4ds88HvEKmE0ik5GwiDbr7EvDDw
+E6UbZtQcIoSTIIZDgYqVFfR2DAho3wXJRsOXh433lEJ8X7cCDzrngFbQnlKrpwML
+Xgm0SIUc+Nf5poMM3rfLFK77t/ob4w+5PwRKcoSniyAxrHd6bwykYA8Vuydv
+-----END RSA PRIVATE KEY-----
diff --git a/mkbootimg.py b/mkbootimg.py
index e0b0839..ec29581 100755
--- a/mkbootimg.py
+++ b/mkbootimg.py
@@ -26,9 +26,10 @@
 import collections
 import os
 import re
-import subprocess
 import tempfile
 
+from gki.generate_gki_certificate import generate_gki_certificate
+
 # Constant and structure definition is in
 # system/tools/mkbootimg/include/bootimg/bootimg.h
 BOOT_MAGIC = 'ANDROID!'
@@ -104,6 +105,12 @@
     return dtbo_offset
 
 
+def should_add_legacy_gki_boot_signature(args):
+    if args.gki_signing_key and args.gki_signing_algorithm:
+        return True
+    return False
+
+
 def write_header_v3_and_above(args):
     if args.header_version > 3:
         boot_header_size = BOOT_IMAGE_HEADER_V4_SIZE
@@ -126,14 +133,14 @@
                            args.cmdline))
     if args.header_version >= 4:
         # The signature used to verify boot image v4.
-        args.output.write(pack('I', BOOT_IMAGE_V4_SIGNATURE_SIZE))
+        boot_signature_size = 0
+        if should_add_legacy_gki_boot_signature(args):
+            boot_signature_size = BOOT_IMAGE_V4_SIGNATURE_SIZE
+        args.output.write(pack('I', boot_signature_size))
     pad_file(args.output, BOOT_IMAGE_HEADER_V3_PAGESIZE)
 
 
 def write_vendor_boot_header(args):
-    if filesize(args.dtb) == 0:
-        raise ValueError('DTB image must not be empty.')
-
     if args.header_version > 3:
         vendor_ramdisk_size = args.vendor_ramdisk_total_size
         vendor_boot_header_size = VENDOR_BOOT_IMAGE_HEADER_V4_SIZE
@@ -535,14 +542,6 @@
                         help='boot image header version')
     parser.add_argument('-o', '--output', type=FileType('wb'),
                         help='output file name')
-    parser.add_argument('--gki_signing_algorithm',
-                        help='GKI signing algorithm to use')
-    parser.add_argument('--gki_signing_key',
-                        help='path to RSA private key file')
-    parser.add_argument('--gki_signing_signature_args',
-                        help='other hash arguments passed to avbtool')
-    parser.add_argument('--gki_signing_avbtool_path',
-                        help='path to avbtool for boot signature generation')
     parser.add_argument('--vendor_boot', type=FileType('wb'),
                         help='vendor boot output file name')
     parser.add_argument('--vendor_ramdisk', type=FileType('rb'),
@@ -550,6 +549,19 @@
     parser.add_argument('--vendor_bootconfig', type=FileType('rb'),
                         help='path to the vendor bootconfig file')
 
+    gki_2_0_signing_args = parser.add_argument_group(
+        '[DEPRECATED] GKI 2.0 signing arguments')
+    gki_2_0_signing_args.add_argument(
+        '--gki_signing_algorithm', help='GKI signing algorithm to use')
+    gki_2_0_signing_args.add_argument(
+        '--gki_signing_key', help='path to RSA private key file')
+    gki_2_0_signing_args.add_argument(
+        '--gki_signing_signature_args', default='',
+        help='other hash arguments passed to avbtool')
+    gki_2_0_signing_args.add_argument(
+        '--gki_signing_avbtool_path', default='avbtool',
+        help='path to avbtool for boot signature generation')
+
     args, extra_args = parser.parse_known_args()
     if args.vendor_boot is not None and args.header_version > 3:
         extra_args = parse_vendor_ramdisk_args(args, extra_args)
@@ -575,50 +587,30 @@
     vbmeta partition) via the Android Verified Boot process, when the
     device boots.
     """
-    args.output.flush()  # Flush the buffer for signature calculation.
-
-    # Appends zeros if the signing key is not specified.
-    if not args.gki_signing_key or not args.gki_signing_algorithm:
-        zeros = b'\x00' * BOOT_IMAGE_V4_SIGNATURE_SIZE
-        args.output.write(zeros)
-        pad_file(args.output, pagesize)
-        return
-
-    avbtool = 'avbtool'  # Used from otatools.zip or Android build env.
-
-    # We need to specify the path of avbtool in build/core/Makefile.
-    # Because avbtool is not guaranteed to be in $PATH there.
-    if args.gki_signing_avbtool_path:
-        avbtool = args.gki_signing_avbtool_path
-
-    # Need to specify a value of --partition_size for avbtool to work.
-    # We use 64 MB below, but avbtool will not resize the boot image to
-    # this size because --do_not_append_vbmeta_image is also specified.
-    avbtool_cmd = [
-        avbtool, 'add_hash_footer',
-        '--partition_name', 'boot',
-        '--partition_size', str(64 * 1024 * 1024),
-        '--image', args.output.name,
-        '--algorithm', args.gki_signing_algorithm,
-        '--key', args.gki_signing_key,
-        '--salt', 'd00df00d']  # TODO: use a hash of kernel/ramdisk as the salt.
-
-    # Additional arguments passed to avbtool.
-    if args.gki_signing_signature_args:
-        avbtool_cmd += args.gki_signing_signature_args.split()
+    # Flush the buffer for signature calculation.
+    args.output.flush()
 
     # Outputs the signed vbmeta to a separate file, then append to boot.img
     # as the boot signature.
     with tempfile.TemporaryDirectory() as temp_out_dir:
         boot_signature_output = os.path.join(temp_out_dir, 'boot_signature')
-        avbtool_cmd += ['--do_not_append_vbmeta_image',
-                        '--output_vbmeta_image', boot_signature_output]
-        subprocess.check_call(avbtool_cmd)
+        generate_gki_certificate(
+            image=args.output.name, avbtool=args.gki_signing_avbtool_path,
+            name='boot', algorithm=args.gki_signing_algorithm,
+            key=args.gki_signing_key, salt='d00df00d',
+            additional_avb_args=args.gki_signing_signature_args.split(),
+            output=boot_signature_output,
+        )
         with open(boot_signature_output, 'rb') as boot_signature:
-            if filesize(boot_signature) > BOOT_IMAGE_V4_SIGNATURE_SIZE:
+            boot_signature_bytes = boot_signature.read()
+            if len(boot_signature_bytes) > BOOT_IMAGE_V4_SIGNATURE_SIZE:
                 raise ValueError(
                     f'boot sigature size is > {BOOT_IMAGE_V4_SIGNATURE_SIZE}')
-            write_padded_file(args.output, boot_signature, pagesize)
+            boot_signature_bytes += b'\x00' * (
+                BOOT_IMAGE_V4_SIGNATURE_SIZE - len(boot_signature_bytes))
+            assert len(boot_signature_bytes) == BOOT_IMAGE_V4_SIGNATURE_SIZE
+            args.output.write(boot_signature_bytes)
+            pad_file(args.output, pagesize)
 
 
 def write_data(args, pagesize):
@@ -630,7 +622,7 @@
         write_padded_file(args.output, args.recovery_dtbo, pagesize)
     if args.header_version == 2:
         write_padded_file(args.output, args.dtb, pagesize)
-    if args.header_version >= 4:
+    if args.header_version >= 4 and should_add_legacy_gki_boot_signature(args):
         add_boot_image_signature(args, pagesize)
 
 
diff --git a/repack_bootimg.py b/repack_bootimg.py
index c320018..93c28f9 100755
--- a/repack_bootimg.py
+++ b/repack_bootimg.py
@@ -128,8 +128,8 @@
                 ['toybox', 'cpio', '-idu'], check=True,
                 input=decompressed_result.stdout, cwd=self._ramdisk_dir)
 
-            print("=== Unpacked ramdisk: '{}' ===".format(
-                self._ramdisk_img))
+            print(f"=== Unpacked ramdisk: '{self._ramdisk_img}' at "
+                  f"'{self._ramdisk_dir}' ===")
         else:
             raise RuntimeError('Failed to decompress ramdisk.')
 
@@ -259,30 +259,22 @@
         subprocess.check_call(mkbootimg_cmd)
         print("=== Repacked boot image: '{}' ===".format(self._bootimg))
 
-    def add_files(self, src_dir, files):
-        """Copy files from the src_dir into current ramdisk.
+    def add_files(self, copy_pairs):
+        """Copy files specified by copy_pairs into current ramdisk.
 
         Args:
-            src_dir: a source dir containing the files to copy from.
-            files: a list of files or src_file:dst_file pairs to copy from
-              src_dir to the current ramdisk.
+            copy_pairs: a list of (src_pathname, dst_file) pairs.
         """
         # Creates missing parent dirs with 0o755.
         original_mask = os.umask(0o022)
-        for f in files:
-            if ':' in f:
-                src_file = os.path.join(src_dir, f.split(':')[0])
-                dst_file = os.path.join(self.ramdisk_dir, f.split(':')[1])
-            else:
-                src_file = os.path.join(src_dir, f)
-                dst_file = os.path.join(self.ramdisk_dir, f)
-
-            dst_dir = os.path.dirname(dst_file)
+        for src_pathname, dst_file in copy_pairs:
+            dst_pathname = os.path.join(self.ramdisk_dir, dst_file)
+            dst_dir = os.path.dirname(dst_pathname)
             if not os.path.exists(dst_dir):
                 print("Creating dir '{}'".format(dst_dir))
                 os.makedirs(dst_dir, 0o755)
-            print("Copying file '{}' into '{}'".format(src_file, dst_file))
-            shutil.copy2(src_file, dst_file)
+            print(f"Copying file '{src_pathname}' to '{dst_pathname}'")
+            shutil.copy2(src_pathname, dst_pathname, follow_symlinks=False)
         os.umask(original_mask)
 
     @property
@@ -294,33 +286,34 @@
 def _get_repack_usage():
     return """Usage examples:
 
-  * --ramdisk_add
+  * --ramdisk_add SRC_FILE:DST_FILE
 
-    Specifies a list of files or src_file:dst_file pairs to copy from
-    --src_bootimg's ramdisk into --dst_bootimg's ramdisk.
+    If --local is given, copy SRC_FILE from the local filesystem to DST_FILE in
+    the ramdisk of --dst_bootimg.
+    If --src_bootimg is specified, copy SRC_FILE from the ramdisk of
+    --src_bootimg to DST_FILE in the ramdisk of --dst_bootimg.
 
-    $ repack_bootimg \\
+    Copies a local file 'userdebug_plat_sepolicy.cil' into the ramdisk of
+    --dst_bootimg, and then rebuild --dst_bootimg:
+
+    $ %(prog)s \\
+        --local --dst_bootimg vendor_boot-debug.img \\
+        --ramdisk_add userdebug_plat_sepolicy.cil:userdebug_plat_sepolicy.cil
+
+    Copies 'first_stage_ramdisk/userdebug_plat_sepolicy.cil' from the ramdisk
+    of --src_bootimg to 'userdebug_plat_sepolicy.cil' in the ramdisk of
+    --dst_bootimg, and then rebuild --dst_bootimg:
+
+    $ %(prog)s \\
         --src_bootimg boot-debug-5.4.img --dst_bootimg vendor_boot-debug.img \\
         --ramdisk_add first_stage_ramdisk/userdebug_plat_sepolicy.cil:userdebug_plat_sepolicy.cil
 
-    The above command copies '/first_stage_ramdisk/userdebug_plat_sepolicy.cil'
-    from --src_bootimg's ramdisk to '/userdebug_plat_sepolicy.cil' of
-    --dst_bootimg's ramdisk, then repacks the --dst_bootimg.
+    This option can be specified multiple times to copy multiple files:
 
-    $ repack_bootimg \\
-        --src_bootimg boot-debug-5.4.img --dst_bootimg vendor_boot-debug.img \\
-        --ramdisk_add first_stage_ramdisk/userdebug_plat_sepolicy.cil
-
-    This is similar to the previous example, but the source file path and
-    destination file path are the same:
-        '/first_stage_ramdisk/userdebug_plat_sepolicy.cil'.
-
-    We can also combine both usage together with a list of copy instructions.
-    For example:
-
-    $ repack_bootimg \\
-        --src_bootimg boot-debug-5.4.img --dst_bootimg vendor_boot-debug.img \\
-        --ramdisk_add file1 file2:/subdir/file2 file3
+    $ %(prog)s \\
+        --local --dst_bootimg vendor_boot-debug.img \\
+        --ramdisk_add file1:path/in/dst_bootimg/file1 \\
+        --ramdisk_add file2:path/in/dst_bootimg/file2
 """
 
 
@@ -328,34 +321,73 @@
     """Parse command-line options."""
     parser = argparse.ArgumentParser(
         formatter_class=argparse.RawDescriptionHelpFormatter,
-        description='Repacks boot, recovery or vendor_boot image by importing'
+        description='Repacks boot, recovery or vendor_boot image by importing '
                     'ramdisk files from --src_bootimg to --dst_bootimg.',
         epilog=_get_repack_usage(),
     )
 
-    parser.add_argument(
+    src_group = parser.add_mutually_exclusive_group(required=True)
+    src_group.add_argument(
         '--src_bootimg', help='filename to source boot image',
-        type=str, required=True)
+        type=BootImage)
+    src_group.add_argument(
+        '--local', help='use local files as repack source',
+        action='store_true')
+
     parser.add_argument(
         '--dst_bootimg', help='filename to destination boot image',
-        type=str, required=True)
+        type=BootImage, required=True)
     parser.add_argument(
-        '--ramdisk_add', nargs='+',
-        help='a list of files or src_file:dst_file pairs to add into '
-             'the ramdisk',
-        default=['userdebug_plat_sepolicy.cil']
-    )
+        '--ramdisk_add', metavar='SRC_FILE:DST_FILE',
+        help='a copy pair to copy into the ramdisk of --dst_bootimg',
+        action='extend', nargs='+', required=True)
 
-    return parser.parse_args()
+    args = parser.parse_args()
+
+    # Parse args.ramdisk_add to a list of copy pairs.
+    if args.src_bootimg:
+        args.ramdisk_add = [
+            _parse_ramdisk_copy_pair(p, args.src_bootimg.ramdisk_dir)
+            for p in args.ramdisk_add
+        ]
+    else:
+        # Repack from local files.
+        args.ramdisk_add = [
+            _parse_ramdisk_copy_pair(p) for p in args.ramdisk_add
+        ]
+
+    return args
+
+
+def _parse_ramdisk_copy_pair(pair, src_ramdisk_dir=None):
+    """Parse a ramdisk copy pair argument."""
+    if ':' in pair:
+        src_file, dst_file = pair.split(':', maxsplit=1)
+    else:
+        src_file = dst_file = pair
+
+    # os.path.join() only works on relative path components.
+    # If a component is an absolute path, all previous components are thrown
+    # away and joining continues from the absolute path component.
+    # So make sure the file name is not absolute before calling os.path.join().
+    if src_ramdisk_dir:
+        if os.path.isabs(src_file):
+            raise ValueError('file name cannot be absolute when repacking from '
+                             'a ramdisk: ' + src_file)
+        src_pathname = os.path.join(src_ramdisk_dir, src_file)
+    else:
+        src_pathname = src_file
+    if os.path.isabs(dst_file):
+        raise ValueError('destination file name cannot be absolute: ' +
+                         dst_file)
+    return (src_pathname, dst_file)
 
 
 def main():
     """Parse arguments and repack boot image."""
     args = _parse_args()
-    src_bootimg = BootImage(args.src_bootimg)
-    dst_bootimg = BootImage(args.dst_bootimg)
-    dst_bootimg.add_files(src_bootimg.ramdisk_dir, args.ramdisk_add)
-    dst_bootimg.repack_bootimg()
+    args.dst_bootimg.add_files(args.ramdisk_add)
+    args.dst_bootimg.repack_bootimg()
 
 
 if __name__ == '__main__':
diff --git a/tests/mkbootimg_test.py b/tests/mkbootimg_test.py
index ae5cf6b..e691e30 100644
--- a/tests/mkbootimg_test.py
+++ b/tests/mkbootimg_test.py
@@ -34,8 +34,6 @@
 VENDOR_BOOT_ARGS_OFFSET = 28
 VENDOR_BOOT_ARGS_SIZE = 2048
 
-BOOT_IMAGE_V4_SIGNATURE_SIZE = 4096
-
 TEST_KERNEL_CMDLINE = (
     'printk.devkmsg=on firmware_class.path=/vendor/etc/ init=/init '
     'kfence.sample_interval=500 loop.max_part=7 bootconfig'
@@ -86,7 +84,7 @@
         # C0103: invalid-name for maxDiff.
         self.maxDiff = None  # pylint: disable=C0103
 
-    def _test_boot_image_v4_signature(self, avbtool_path):
+    def _test_legacy_boot_image_v4_signature(self, avbtool_path):
         """Tests the boot_signature in boot.img v4."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
             boot_img = os.path.join(temp_out_dir, 'boot.img')
@@ -162,15 +160,16 @@
 
             self.assertEqual(result.stdout, expected_boot_signature_info)
 
-    def test_boot_image_v4_signature_without_avbtool_path(self):
+    def test_legacy_boot_image_v4_signature_without_avbtool_path(self):
         """Boot signature generation without --gki_signing_avbtool_path."""
-        self._test_boot_image_v4_signature(avbtool_path=None)
+        self._test_legacy_boot_image_v4_signature(avbtool_path=None)
 
-    def test_boot_image_v4_signature_with_avbtool_path(self):
+    def test_legacy_boot_image_v4_signature_with_avbtool_path(self):
         """Boot signature generation with --gki_signing_avbtool_path."""
-        self._test_boot_image_v4_signature(avbtool_path=self._avbtool_path)
+        self._test_legacy_boot_image_v4_signature(
+            avbtool_path=self._avbtool_path)
 
-    def test_boot_image_v4_signature_exceed_size(self):
+    def test_legacy_boot_image_v4_signature_exceed_size(self):
         """Tests the boot signature size exceeded in a boot image version 4."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
             boot_img = os.path.join(temp_out_dir, 'boot.img')
@@ -205,7 +204,7 @@
                 self.assertIn('ValueError: boot sigature size is > 4096',
                               e.stderr)
 
-    def test_boot_image_v4_signature_zeros(self):
+    def test_boot_image_v4_signature_empty(self):
         """Tests no boot signature in a boot image version 4."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
             boot_img = os.path.join(temp_out_dir, 'boot.img')
@@ -214,8 +213,6 @@
             ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
                                          0x1000)
 
-            # The boot signature will be zeros if no
-            # --gki_signing_[algorithm|key] is provided.
             mkbootimg_cmds = [
                 'mkbootimg',
                 '--header_version', '4',
@@ -235,11 +232,10 @@
             subprocess.run(mkbootimg_cmds, check=True)
             subprocess.run(unpack_bootimg_cmds, check=True)
 
-            boot_signature = os.path.join(
-                temp_out_dir, 'out', 'boot_signature')
-            with open(boot_signature) as f:
-                zeros = '\x00' * BOOT_IMAGE_V4_SIGNATURE_SIZE
-                self.assertEqual(f.read(), zeros)
+            # The boot signature will be empty if no
+            # --gki_signing_[algorithm|key] is provided.
+            boot_signature = os.path.join(temp_out_dir, 'out', 'boot_signature')
+            self.assertFalse(os.path.exists(boot_signature))
 
     def test_vendor_boot_v4(self):
         """Tests vendor_boot version 4."""
@@ -418,6 +414,45 @@
                 filecmp.cmp(vendor_boot_img, vendor_boot_img_reconstructed),
                 'reconstructed vendor_boot image differ from the original')
 
+    def test_unpack_boot_image_v4(self):
+        """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            boot_img = os.path.join(temp_out_dir, 'boot.img')
+            boot_img_reconstructed = os.path.join(
+                temp_out_dir, 'boot.img.reconstructed')
+            kernel = generate_test_file(os.path.join(temp_out_dir, 'kernel'),
+                                        0x1000)
+            ramdisk = generate_test_file(os.path.join(temp_out_dir, 'ramdisk'),
+                                         0x1000)
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--header_version', '4',
+                '--kernel', kernel,
+                '--ramdisk', ramdisk,
+                '--cmdline', TEST_KERNEL_CMDLINE,
+                '--output', boot_img,
+            ]
+            unpack_bootimg_cmds = [
+                'unpack_bootimg',
+                '--boot_img', boot_img,
+                '--out', os.path.join(temp_out_dir, 'out'),
+                '--format=mkbootimg',
+            ]
+
+            subprocess.run(mkbootimg_cmds, check=True)
+            result = subprocess.run(unpack_bootimg_cmds, check=True,
+                                    capture_output=True, encoding='utf-8')
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--out', boot_img_reconstructed,
+            ]
+            mkbootimg_cmds.extend(shlex.split(result.stdout))
+
+            subprocess.run(mkbootimg_cmds, check=True)
+            self.assertTrue(
+                filecmp.cmp(boot_img, boot_img_reconstructed),
+                'reconstructed boot image differ from the original')
+
     def test_unpack_boot_image_v3(self):
         """Tests that mkbootimg(unpack_bootimg(image)) is an identity."""
         with tempfile.TemporaryDirectory() as temp_out_dir:
@@ -725,6 +760,80 @@
             self.assertEqual(raw_vendor_cmdline,
                              vendor_cmdline.encode() + b'\x00')
 
+    def test_vendor_boot_v4_without_dtb(self):
+        """Tests building vendor_boot version 4 without dtb image."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            vendor_boot_img = os.path.join(temp_out_dir, 'vendor_boot.img')
+            ramdisk = generate_test_file(
+                os.path.join(temp_out_dir, 'ramdisk'), 0x1000)
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--header_version', '4',
+                '--vendor_boot', vendor_boot_img,
+                '--vendor_ramdisk', ramdisk,
+            ]
+            unpack_bootimg_cmds = [
+                'unpack_bootimg',
+                '--boot_img', vendor_boot_img,
+                '--out', os.path.join(temp_out_dir, 'out'),
+            ]
+            expected_output = [
+                'boot magic: VNDRBOOT',
+                'vendor boot image header version: 4',
+                'dtb size: 0',
+            ]
+
+            subprocess.run(mkbootimg_cmds, check=True)
+            result = subprocess.run(unpack_bootimg_cmds, check=True,
+                                    capture_output=True, encoding='utf-8')
+            output = [line.strip() for line in result.stdout.splitlines()]
+            if not subsequence_of(expected_output, output):
+                msg = '\n'.join([
+                    'Unexpected unpack_bootimg output:',
+                    'Expected:',
+                    ' ' + '\n '.join(expected_output),
+                    '',
+                    'Actual:',
+                    ' ' + '\n '.join(output),
+                ])
+                self.fail(msg)
+
+    def test_unpack_vendor_boot_image_v4_without_dtb(self):
+        """Tests that mkbootimg(unpack_bootimg(image)) is an identity when no dtb image."""
+        with tempfile.TemporaryDirectory() as temp_out_dir:
+            vendor_boot_img = os.path.join(temp_out_dir, 'vendor_boot.img')
+            vendor_boot_img_reconstructed = os.path.join(
+                temp_out_dir, 'vendor_boot.img.reconstructed')
+            ramdisk = generate_test_file(
+                os.path.join(temp_out_dir, 'ramdisk'), 0x121212)
+
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--header_version', '4',
+                '--vendor_boot', vendor_boot_img,
+                '--vendor_ramdisk', ramdisk,
+            ]
+            unpack_bootimg_cmds = [
+                'unpack_bootimg',
+                '--boot_img', vendor_boot_img,
+                '--out', os.path.join(temp_out_dir, 'out'),
+                '--format=mkbootimg',
+            ]
+            subprocess.run(mkbootimg_cmds, check=True)
+            result = subprocess.run(unpack_bootimg_cmds, check=True,
+                                    capture_output=True, encoding='utf-8')
+            mkbootimg_cmds = [
+                'mkbootimg',
+                '--vendor_boot', vendor_boot_img_reconstructed,
+            ]
+            unpack_format_args = shlex.split(result.stdout)
+            mkbootimg_cmds.extend(unpack_format_args)
+
+            subprocess.run(mkbootimg_cmds, check=True)
+            self.assertTrue(
+                filecmp.cmp(vendor_boot_img, vendor_boot_img_reconstructed),
+                'reconstructed vendor_boot image differ from the original')
+
 
 # I don't know how, but we need both the logger configuration and verbosity
 # level > 2 to make atest work. And yes this line needs to be at the very top
diff --git a/unpack_bootimg.py b/unpack_bootimg.py
index 2b176e5..462190f 100755
--- a/unpack_bootimg.py
+++ b/unpack_bootimg.py
@@ -19,7 +19,7 @@
 Extracts the kernel, ramdisk, second bootloader, dtb and recovery dtbo images.
 """
 
-from argparse import ArgumentParser, FileType, RawDescriptionHelpFormatter
+from argparse import ArgumentParser, RawDescriptionHelpFormatter
 from struct import unpack
 import os
 import shlex
@@ -53,17 +53,21 @@
 
 
 def format_os_version(os_version):
+    if os_version == 0:
+        return None
     a = os_version >> 14
     b = os_version >> 7 & ((1<<7) - 1)
     c = os_version & ((1<<7) - 1)
-    return '{}.{}.{}'.format(a, b, c)
+    return f'{a}.{b}.{c}'
 
 
 def format_os_patch_level(os_patch_level):
+    if os_patch_level == 0:
+        return None
     y = os_patch_level >> 4
     y += 2000
     m = os_patch_level & ((1<<4) - 1)
-    return '{:04d}-{:02d}'.format(y, m)
+    return f'{y:04d}-{m:02d}'
 
 
 def decode_os_version_patch_level(os_version_patch_level):
@@ -130,8 +134,10 @@
     def format_mkbootimg_argument(self):
         args = []
         args.extend(['--header_version', str(self.header_version)])
-        args.extend(['--os_version', self.os_version])
-        args.extend(['--os_patch_level', self.os_patch_level])
+        if self.os_version:
+            args.extend(['--os_version', self.os_version])
+        if self.os_patch_level:
+            args.extend(['--os_patch_level', self.os_patch_level])
 
         args.extend(['--kernel', os.path.join(self.image_dir, 'kernel')])
         args.extend(['--ramdisk', os.path.join(self.image_dir, 'ramdisk')])
@@ -175,12 +181,12 @@
         return args
 
 
-def unpack_boot_image(args):
+def unpack_boot_image(boot_img, output_dir):
     """extracts kernel, ramdisk, second bootloader and recovery dtbo"""
     info = BootImageInfoFormatter()
-    info.boot_magic = unpack('8s', args.boot_img.read(8))[0].decode()
+    info.boot_magic = unpack('8s', boot_img.read(8))[0].decode()
 
-    kernel_ramdisk_second_info = unpack('9I', args.boot_img.read(9 * 4))
+    kernel_ramdisk_second_info = unpack('9I', boot_img.read(9 * 4))
     # header_version is always at [8] regardless of the value of header_version.
     info.header_version = kernel_ramdisk_second_info[8]
 
@@ -193,7 +199,7 @@
         info.second_load_address = kernel_ramdisk_second_info[5]
         info.tags_load_address = kernel_ramdisk_second_info[6]
         info.page_size = kernel_ramdisk_second_info[7]
-        os_version_patch_level = unpack('I', args.boot_img.read(1 * 4))[0]
+        os_version_patch_level = unpack('I', boot_img.read(1 * 4))[0]
     else:
         info.kernel_size = kernel_ramdisk_second_info[0]
         info.ramdisk_size = kernel_ramdisk_second_info[1]
@@ -206,31 +212,31 @@
 
     if info.header_version < 3:
         info.product_name = cstr(unpack('16s',
-                                        args.boot_img.read(16))[0].decode())
-        info.cmdline = cstr(unpack('512s', args.boot_img.read(512))[0].decode())
-        args.boot_img.read(32)  # ignore SHA
+                                        boot_img.read(16))[0].decode())
+        info.cmdline = cstr(unpack('512s', boot_img.read(512))[0].decode())
+        boot_img.read(32)  # ignore SHA
         info.extra_cmdline = cstr(unpack('1024s',
-                                         args.boot_img.read(1024))[0].decode())
+                                         boot_img.read(1024))[0].decode())
     else:
         info.cmdline = cstr(unpack('1536s',
-                                   args.boot_img.read(1536))[0].decode())
+                                   boot_img.read(1536))[0].decode())
 
     if info.header_version in {1, 2}:
-        info.recovery_dtbo_size = unpack('I', args.boot_img.read(1 * 4))[0]
-        info.recovery_dtbo_offset = unpack('Q', args.boot_img.read(8))[0]
-        info.boot_header_size = unpack('I', args.boot_img.read(4))[0]
+        info.recovery_dtbo_size = unpack('I', boot_img.read(1 * 4))[0]
+        info.recovery_dtbo_offset = unpack('Q', boot_img.read(8))[0]
+        info.boot_header_size = unpack('I', boot_img.read(4))[0]
     else:
         info.recovery_dtbo_size = 0
 
     if info.header_version == 2:
-        info.dtb_size = unpack('I', args.boot_img.read(4))[0]
-        info.dtb_load_address = unpack('Q', args.boot_img.read(8))[0]
+        info.dtb_size = unpack('I', boot_img.read(4))[0]
+        info.dtb_load_address = unpack('Q', boot_img.read(8))[0]
     else:
         info.dtb_size = 0
         info.dtb_load_address = 0
 
     if info.header_version >= 4:
-        info.boot_signature_size = unpack('I', args.boot_img.read(4))[0]
+        info.boot_signature_size = unpack('I', boot_img.read(4))[0]
     else:
         info.boot_signature_size = 0
 
@@ -278,10 +284,10 @@
         image_info_list.append((boot_signature_offset, info.boot_signature_size,
                                 'boot_signature'))
 
-    create_out_dir(args.out)
+    create_out_dir(output_dir)
     for offset, size, name in image_info_list:
-        extract_image(offset, size, args.boot_img, os.path.join(args.out, name))
-    info.image_dir = args.out
+        extract_image(offset, size, boot_img, os.path.join(output_dir, name))
+    info.image_dir = output_dir
 
     return info
 
@@ -347,7 +353,8 @@
         args.extend(['--vendor_cmdline', self.cmdline])
         args.extend(['--board', self.product_name])
 
-        args.extend(['--dtb', os.path.join(self.image_dir, 'dtb')])
+        if self.dtb_size > 0:
+            args.extend(['--dtb', os.path.join(self.image_dir, 'dtb')])
 
         if self.header_version > 3:
             args.extend(['--vendor_bootconfig',
@@ -371,20 +378,20 @@
         return args
 
 
-def unpack_vendor_boot_image(args):
+def unpack_vendor_boot_image(boot_img, output_dir):
     info = VendorBootImageInfoFormatter()
-    info.boot_magic = unpack('8s', args.boot_img.read(8))[0].decode()
-    info.header_version = unpack('I', args.boot_img.read(4))[0]
-    info.page_size = unpack('I', args.boot_img.read(4))[0]
-    info.kernel_load_address = unpack('I', args.boot_img.read(4))[0]
-    info.ramdisk_load_address = unpack('I', args.boot_img.read(4))[0]
-    info.vendor_ramdisk_size = unpack('I', args.boot_img.read(4))[0]
-    info.cmdline = cstr(unpack('2048s', args.boot_img.read(2048))[0].decode())
-    info.tags_load_address = unpack('I', args.boot_img.read(4))[0]
-    info.product_name = cstr(unpack('16s', args.boot_img.read(16))[0].decode())
-    info.header_size = unpack('I', args.boot_img.read(4))[0]
-    info.dtb_size = unpack('I', args.boot_img.read(4))[0]
-    info.dtb_load_address = unpack('Q', args.boot_img.read(8))[0]
+    info.boot_magic = unpack('8s', boot_img.read(8))[0].decode()
+    info.header_version = unpack('I', boot_img.read(4))[0]
+    info.page_size = unpack('I', boot_img.read(4))[0]
+    info.kernel_load_address = unpack('I', boot_img.read(4))[0]
+    info.ramdisk_load_address = unpack('I', boot_img.read(4))[0]
+    info.vendor_ramdisk_size = unpack('I', boot_img.read(4))[0]
+    info.cmdline = cstr(unpack('2048s', boot_img.read(2048))[0].decode())
+    info.tags_load_address = unpack('I', boot_img.read(4))[0]
+    info.product_name = cstr(unpack('16s', boot_img.read(16))[0].decode())
+    info.header_size = unpack('I', boot_img.read(4))[0]
+    info.dtb_size = unpack('I', boot_img.read(4))[0]
+    info.dtb_load_address = unpack('Q', boot_img.read(8))[0]
 
     # Convenient shorthand.
     page_size = info.page_size
@@ -397,11 +404,14 @@
     ramdisk_offset_base = page_size * num_boot_header_pages
     image_info_list = []
 
+    image_info_list.append(
+        (ramdisk_offset_base, info.vendor_ramdisk_size, 'vendor_ramdisk'))
+
     if info.header_version > 3:
-        info.vendor_ramdisk_table_size = unpack('I', args.boot_img.read(4))[0]
-        vendor_ramdisk_table_entry_num = unpack('I', args.boot_img.read(4))[0]
-        vendor_ramdisk_table_entry_size = unpack('I', args.boot_img.read(4))[0]
-        info.vendor_bootconfig_size = unpack('I', args.boot_img.read(4))[0]
+        info.vendor_ramdisk_table_size = unpack('I', boot_img.read(4))[0]
+        vendor_ramdisk_table_entry_num = unpack('I', boot_img.read(4))[0]
+        vendor_ramdisk_table_entry_size = unpack('I', boot_img.read(4))[0]
+        info.vendor_bootconfig_size = unpack('I', boot_img.read(4))[0]
         num_vendor_ramdisk_table_pages = get_number_of_pages(
             info.vendor_ramdisk_table_size, page_size)
         vendor_ramdisk_table_offset = page_size * (
@@ -412,16 +422,16 @@
         for idx in range(vendor_ramdisk_table_entry_num):
             entry_offset = vendor_ramdisk_table_offset + (
                 vendor_ramdisk_table_entry_size * idx)
-            args.boot_img.seek(entry_offset)
-            ramdisk_size = unpack('I', args.boot_img.read(4))[0]
-            ramdisk_offset = unpack('I', args.boot_img.read(4))[0]
-            ramdisk_type = unpack('I', args.boot_img.read(4))[0]
+            boot_img.seek(entry_offset)
+            ramdisk_size = unpack('I', boot_img.read(4))[0]
+            ramdisk_offset = unpack('I', boot_img.read(4))[0]
+            ramdisk_type = unpack('I', boot_img.read(4))[0]
             ramdisk_name = cstr(unpack(
                 f'{VENDOR_RAMDISK_NAME_SIZE}s',
-                args.boot_img.read(VENDOR_RAMDISK_NAME_SIZE))[0].decode())
+                boot_img.read(VENDOR_RAMDISK_NAME_SIZE))[0].decode())
             board_id = unpack(
                 f'{VENDOR_RAMDISK_TABLE_ENTRY_BOARD_ID_SIZE}I',
-                args.boot_img.read(
+                boot_img.read(
                     4 * VENDOR_RAMDISK_TABLE_ENTRY_BOARD_ID_SIZE))
             output_ramdisk_name = f'vendor_ramdisk{idx:02}'
 
@@ -439,22 +449,20 @@
             + num_vendor_ramdisk_table_pages)
         image_info_list.append((bootconfig_offset, info.vendor_bootconfig_size,
             'bootconfig'))
-    else:
-        image_info_list.append(
-            (ramdisk_offset_base, info.vendor_ramdisk_size, 'vendor_ramdisk'))
 
     dtb_offset = page_size * (num_boot_header_pages + num_boot_ramdisk_pages
                              ) # header + vendor_ramdisk
-    image_info_list.append((dtb_offset, info.dtb_size, 'dtb'))
+    if info.dtb_size > 0:
+        image_info_list.append((dtb_offset, info.dtb_size, 'dtb'))
 
-    create_out_dir(args.out)
+    create_out_dir(output_dir)
     for offset, size, name in image_info_list:
-        extract_image(offset, size, args.boot_img, os.path.join(args.out, name))
-    info.image_dir = args.out
+        extract_image(offset, size, boot_img, os.path.join(output_dir, name))
+    info.image_dir = output_dir
 
     if info.header_version > 3:
         vendor_ramdisk_by_name_dir = os.path.join(
-            args.out, 'vendor-ramdisk-by-name')
+            output_dir, 'vendor-ramdisk-by-name')
         create_out_dir(vendor_ramdisk_by_name_dir)
         for src, dst in vendor_ramdisk_symlinks:
             src_pathname = os.path.join('..', src)
@@ -467,19 +475,26 @@
     return info
 
 
-def unpack_image(args):
-    boot_magic = unpack('8s', args.boot_img.read(8))[0].decode()
-    args.boot_img.seek(0)
-    if boot_magic == 'ANDROID!':
-        info = unpack_boot_image(args)
-    elif boot_magic == 'VNDRBOOT':
-        info = unpack_vendor_boot_image(args)
-    else:
-        raise ValueError(f'Not an Android boot image, magic: {boot_magic}')
+def unpack_bootimg(boot_img, output_dir):
+    """Unpacks the |boot_img| to |output_dir|, and returns the 'info' object."""
+    with open(boot_img, 'rb') as image_file:
+        boot_magic = unpack('8s', image_file.read(8))[0].decode()
+        image_file.seek(0)
+        if boot_magic == 'ANDROID!':
+            info = unpack_boot_image(image_file, output_dir)
+        elif boot_magic == 'VNDRBOOT':
+            info = unpack_vendor_boot_image(image_file, output_dir)
+        else:
+            raise ValueError(f'Not an Android boot image, magic: {boot_magic}')
 
-    if args.format == 'mkbootimg':
+    return info
+
+
+def print_bootimg_info(info, output_format, null_separator):
+    """Format and print boot image info."""
+    if output_format == 'mkbootimg':
         mkbootimg_args = info.format_mkbootimg_argument()
-        if args.null:
+        if null_separator:
             print('\0'.join(mkbootimg_args) + '\0', end='')
         else:
             print(shlex.join(mkbootimg_args))
@@ -525,7 +540,7 @@
         description='Unpacks boot, recovery or vendor_boot image.',
         epilog=get_unpack_usage(),
     )
-    parser.add_argument('--boot_img', type=FileType('rb'), required=True,
+    parser.add_argument('--boot_img', required=True,
                         help='path to the boot, recovery or vendor_boot image')
     parser.add_argument('--out', default='out',
                         help='output directory of the unpacked images')
@@ -540,7 +555,8 @@
 def main():
     """parse arguments and unpack boot image"""
     args = parse_cmdline()
-    unpack_image(args)
+    info = unpack_bootimg(args.boot_img, args.out)
+    print_bootimg_info(info, args.format, args.null)
 
 
 if __name__ == '__main__':