sign_official_build: Port google3 changes for COS

As part of the Cusky signing design, we carry custom changes to the
sign_official_build.sh script in a fork of vboot_reference repository on
google3. To move the signing behavior from the BC to the BE, we will need
to port those changes from google3 to here.

To make rebasing this change when do we do a ChromeOS rebase easier, we
fork the changes that were made in google3 into a new file,
sign_official_cos_build.sh

The custom changes being ported here are the ones made in cl/316911016,
cl/317144168, cl/318860928, cl/318866697 and cl/318866697.

BUG=b/170128847
TEST=presubmit; merged signer container with BE container, patched these
changes in and ran the script

Change-Id: I75af70e0b3603d6fdc4ce05ed8383ab8a66e05fb
Reviewed-on: https://cos-review.googlesource.com/c/third_party/platform/vboot_reference/+/18671
Reviewed-by: Robert Kolchmeyer <rkolchmeyer@google.com>
Tested-by: Dexter Rivera <riverade@google.com>
Reviewed-on: https://cos-review.googlesource.com/c/third_party/platform/vboot_reference/+/21903
Tested-by: Cusky QA Presubmit Bot <presubmit@cos-infra-prototype.iam.gserviceaccount.com>
diff --git a/scripts/image_signing/sign_official_cos_build.sh b/scripts/image_signing/sign_official_cos_build.sh
new file mode 100755
index 0000000..1e859c9
--- /dev/null
+++ b/scripts/image_signing/sign_official_cos_build.sh
@@ -0,0 +1,731 @@
+#!/bin/bash
+
+# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+# Sign the final build image using the "official" keys.
+#
+# Prerequisite tools needed in the system path:
+#
+#  futility (from src/platform/vboot_reference)
+#  vbutil_kernel (from src/platform/vboot_reference)
+#  vbutil_key (from src/platform/vboot_reference)
+#  cgpt (from src/platform/vboot_reference)
+#  dump_kernel_config (from src/platform/vboot_reference)
+#  verity (from src/platform/verity)
+#  load_kernel_test (from src/platform/vboot_reference)
+#  dumpe2fs
+#  sha1sum
+#  kms_signer, if KEY_ORIGIN == kms.
+
+# Load common constants and variables.
+. "$(dirname "$0")/common.sh"
+
+# Print usage string
+usage() {
+  cat <<EOF
+Usage: $PROG <type> <key_origin> input_image /path/to/keys/dir [output_image] [version_file]
+where <type> is one of:
+             base (sign a base image, similar to an SSD image)
+             update_payload (sign a delta update hash)
+key_origin: "local" for local keys, or "kms" for Cloud KMS keys.
+output_image: File name of the signed output image
+version_file: File name of where to read the kernel and firmware versions.
+
+If you are signing an image, you must specify an [output_image] and
+optionally, a [version_file].
+
+EOF
+  if [[ $# -gt 0 ]]; then
+    error "$*"
+    exit 1
+  fi
+  exit 0
+}
+
+# Verify we have as many arguments as we expect, else show usage & quit.
+# Usage:
+#  check_argc <number args> <exact number>
+#  check_argc <number args> <lower bound> <upper bound>
+check_argc() {
+  case $# in
+  2)
+    if [[ $1 -ne $2 ]]; then
+      usage "command takes exactly $2 args"
+    fi
+    ;;
+  3)
+    if [[ $1 -lt $2 || $1 -gt $3 ]]; then
+      usage "command takes $2 to $3 args"
+    fi
+    ;;
+  *)
+    die "check_argc: incorrect number of arguments"
+  esac
+}
+
+# Abort on errors.
+set -e
+
+# Add to the path since some tools reside here and may not be in the non-root
+# system path.
+PATH=$PATH:/usr/sbin:/sbin
+
+TYPE=$1
+KEY_ORIGIN=$2
+INPUT_IMAGE=$3
+KEY_DIR=$4
+OUTPUT_IMAGE=$5
+VERSION_FILE=$6
+
+FIRMWARE_VERSION=1
+KERNEL_VERSION=1
+
+# Make sure the tools we need are available.
+prereqs=("${FUTILITY}" vbutil_kernel cgpt dump_kernel_config verity)
+prereqs+=(load_kernel_test dumpe2fs sha1sum e2fsck)
+if [[ "${KEY_ORIGIN}" == "kms" ]]; then
+  prereqs+=(kms_signer)
+fi
+for prereq in ${prereqs[@]}; do
+  type -P "${prereq}" &>/dev/null || \
+    die "${prereq} tool not found."
+done
+
+# TODO(gauravsh): These are duplicated from chromeos-setimage. We need
+# to move all signing and rootfs code to one single place where it can be
+# reused. crosbug.com/19543
+
+# get_verity_arg <commandline> <key> -> <value>
+get_verity_arg() {
+  echo "$1" | sed -n "s/.*\b$2=\([^ \"]*\).*/\1/p"
+}
+
+# Get the dmparams parameters from a kernel config.
+get_dmparams_from_config() {
+  local kernel_config=$1
+  echo ${kernel_config} | sed -nre 's/.*dm="([^"]*)".*/\1/p'
+}
+# Get the verity root digest hash from a kernel config command line.
+get_hash_from_config() {
+  local kernel_config=$1
+  local dm_config=$(get_dmparams_from_config "${kernel_config}")
+  local vroot_dev=$(get_dm_slave "${dm_config}" vroot)
+  echo $(get_verity_arg "${vroot_dev}" root_hexdigest)
+}
+
+# Get the slave device and its args
+# get_dm_ags $dm_config [vboot|vroot]
+# Assumes we have only one slave device per device
+get_dm_slave() {
+  local dm=$1
+  local device=$2
+  echo $(echo "${dm}" | sed -nre "s/.*${device}[^,]*,([^,]*).*/\1/p")
+}
+
+# Set the slave device and its args for a device
+# get_dm_ags $dm_config [vboot|vroot] args
+# Assumes we have only one slave device per device
+set_dm_slave() {
+  local dm=$1
+  local device=$2
+  local slave=$3
+  echo $(echo "${dm}" |
+    sed -nre "s#(.*${device}[^,]*,)([^,]*)(.*)#\1${slave}\3#p")
+}
+
+CALCULATED_KERNEL_CONFIG=
+CALCULATED_DM_ARGS=
+# Calculate rootfs hash of an image
+# Args: ROOTFS_IMAGE KERNEL_CONFIG HASH_IMAGE
+#
+# rootfs calculation parameters are grabbed from KERNEL_CONFIG
+#
+# Updated dm-verity arguments (to be replaced in kernel config command line)
+# with the new hash is stored in $CALCULATED_DM_ARGS and the new hash image is
+# written to the file HASH_IMAGE.
+calculate_rootfs_hash() {
+  local rootfs_image=$1
+  local kernel_config=$2
+  local hash_image=$3
+  local dm_config=$(get_dmparams_from_config "${kernel_config}")
+
+  if [ -z "${dm_config}" ]; then
+    warn "Couldn't grab dm_config. Aborting rootfs hash calculation."
+    return 1
+  fi
+  local vroot_dev=$(get_dm_slave "${dm_config}" vroot)
+
+  # Extract the key-value parameters from the kernel command line.
+  local rootfs_sectors=$(get_verity_arg "${vroot_dev}" hashstart)
+  local verity_algorithm=$(get_verity_arg "${vroot_dev}" alg)
+  local root_dev=$(get_verity_arg "${vroot_dev}" payload)
+  local hash_dev=$(get_verity_arg "${vroot_dev}" hashtree)
+  local salt=$(get_verity_arg "${vroot_dev}" salt)
+
+  local salt_arg
+  if [ -n "$salt" ]; then
+    salt_arg="salt=$salt"
+  fi
+
+  # Run the verity tool on the rootfs partition.
+  local slave=$(sudo verity mode=create \
+    alg=${verity_algorithm} \
+    payload="${rootfs_image}" \
+    payload_blocks=$((rootfs_sectors / 8)) \
+    hashtree="${hash_image}" ${salt_arg})
+  # Reconstruct new kernel config command line and replace placeholders.
+  slave="$(echo "${slave}" |
+    sed -s "s|ROOT_DEV|${root_dev}|g;s|HASH_DEV|${hash_dev}|")"
+  CALCULATED_DM_ARGS="$(set_dm_slave "${dm_config}" vroot "${slave}")"
+  CALCULATED_KERNEL_CONFIG="$(echo "${kernel_config}" |
+    sed -e 's#\(.*dm="\)\([^"]*\)\(".*\)'"#\1${CALCULATED_DM_ARGS}\3#g")"
+}
+
+# Re-calculate rootfs hash, update rootfs and kernel command line(s).
+# Args: LOOPDEV KERNEL KERN_A_KEYBLOCK KERN_A_PRIVKEY KERN_B_KEYBLOCK \
+#       KERN_B_PRIVKEY
+#
+# The rootfs is hashed by tool 'verity', and the hash data is stored after the
+# rootfs. A hash of those hash data (also known as final verity hash) may be
+# contained in kernel 2 or kernel 4 command line.
+#
+# This function reads dm-verity configuration from KERNEL, rebuilds the rootfs
+# hash, and then resigns kernel A & B by their keyblock and private key files.
+update_rootfs_hash() {
+  local loopdev="$1"  # Input image.
+  local loop_kern="$2"  # Kernel that contains verity args.
+  local kern_a_keyblock="$3"  # Keyblock file for kernel A.
+  local kern_a_privkey="$4"  # Private key file for kernel A.
+  local kern_b_keyblock="$5"  # Keyblock file for kernel B.
+  local kern_b_privkey="$6"  # Private key file for kernel A.
+  local loop_rootfs="${loopdev}p3"
+
+  # Note even though there are two kernels, there is one place (after rootfs)
+  # for hash data, so we must assume both kernel use same hash algorithm (i.e.,
+  # DM config).
+  info "Updating rootfs hash and updating config for Kernel partitions"
+
+  # If we can't find dm parameters in the kernel config, bail out now.
+  local kernel_config=$(sudo dump_kernel_config "${loop_kern}")
+  local dm_config=$(get_dmparams_from_config "${kernel_config}")
+  if [ -z "${dm_config}" ]; then
+    error "Couldn't grab dm_config from kernel ${loop_kern}"
+    error " (config: ${kernel_config})"
+    return 1
+  fi
+
+  # check and clear need_to_resign tag
+  local rootfs_dir=$(make_temp_dir)
+  sudo mount -o ro "${loop_rootfs}" "${rootfs_dir}"
+  if has_needs_to_be_resigned_tag "${rootfs_dir}"; then
+    # remount as RW
+    sudo mount -o remount,rw "${rootfs_dir}"
+    sudo rm -f "${rootfs_dir}/${TAG_NEEDS_TO_BE_SIGNED}"
+  fi
+  sudo umount "${rootfs_dir}"
+
+  local hash_image=$(make_temp_file)
+
+  # Disable rw mount support prior to hashing.
+  disable_rw_mount "${loop_rootfs}"
+
+  if ! calculate_rootfs_hash "${loop_rootfs}"  "${kernel_config}" \
+    "${hash_image}"; then
+    error "calculate_rootfs_hash failed!"
+    error "Aborting rootfs hash update!"
+    return 1
+  fi
+
+  local rootfs_blocks=$(sudo dumpe2fs "${loop_rootfs}" 2> /dev/null |
+    grep "Block count" |
+    tr -d ' ' |
+    cut -f2 -d:)
+  local rootfs_sectors=$((rootfs_blocks * 8))
+
+  # Overwrite the appended hashes in the rootfs
+  sudo dd if="${hash_image}" of="${loop_rootfs}" bs=512 \
+    seek=${rootfs_sectors} conv=notrunc 2>/dev/null
+
+  # Update kernel command lines
+  local dm_args="${CALCULATED_DM_ARGS}"
+  local temp_config=$(make_temp_file)
+  local kernelpart=
+  local keyblock=
+  local priv_key=
+  local new_kernel_config=
+
+  for kernelpart in 2 4; do
+    loop_kern="${loopdev}p${kernelpart}"
+    if ! new_kernel_config="$(
+         sudo dump_kernel_config "${loop_kern}" 2>/dev/null)" &&
+       [[ "${kernelpart}" == 4 ]]; then
+      # Legacy images don't have partition 4.
+      info "Skipping empty kernel partition 4 (legacy images)."
+      continue
+    fi
+    new_kernel_config="$(echo "${new_kernel_config}" |
+      sed -e 's#\(.*dm="\)\([^"]*\)\(".*\)'"#\1${dm_args}\3#g")"
+    info "New config for kernel partition ${kernelpart} is:"
+    echo "${new_kernel_config}" | tee "${temp_config}"
+    # Re-calculate kernel partition signature and command line.
+    if [[ "$kernelpart" == 2 ]]; then
+      keyblock="${kern_a_keyblock}"
+      priv_key="${kern_a_privkey}"
+    else
+      keyblock="${kern_b_keyblock}"
+      priv_key="${kern_b_privkey}"
+    fi
+    sudo vbutil_kernel --repack "${loop_kern}" \
+      --keyblock ${keyblock} \
+      --signprivate ${priv_key} \
+      --version "${KERNEL_VERSION}" \
+      --oldblob "${loop_kern}" \
+      --config ${temp_config}
+  done
+}
+
+# Update the SSD install-able vblock file on stateful partition.
+# ARGS: Loopdev
+# This is deprecated because all new images should have a SSD boot-able kernel
+# in partition 4. However, the signer needs to be able to sign new & old images
+# (crbug.com/449450#c13) so we will probably never remove this.
+update_stateful_partition_vblock() {
+  local loopdev="$1"
+  local temp_out_vb="$(make_temp_file)"
+
+  local loop_kern="${loopdev}p4"
+  if [[ -z "$(sudo dump_kernel_config "${loop_kern}" 2>/dev/null)" ]]; then
+    info "Building vmlinuz_hd.vblock from legacy image partition 2."
+    loop_kern="${loopdev}p2"
+  fi
+
+  # vblock should always use kernel keyblock.
+  sudo vbutil_kernel --repack "${temp_out_vb}" \
+    --keyblock "${KEY_DIR}/kernel.keyblock" \
+    --signprivate "${KEY_DIR}/kernel_data_key.vbprivk" \
+    --oldblob "${loop_kern}" \
+    --vblockonly
+
+  # Copy the installer vblock to the stateful partition.
+  local stateful_dir=$(make_temp_dir)
+  sudo mount "${loopdev}p1" "${stateful_dir}"
+  sudo cp ${temp_out_vb} ${stateful_dir}/vmlinuz_hd.vblock
+  sudo umount "${stateful_dir}"
+}
+
+# Do a sanity check on the image's rootfs
+# ARGS: Image
+verify_image_rootfs() {
+  local rootfs=$1
+  # This flips the read-only compatibility flag, so that e2fsck does not
+  # complain about unknown file system capabilities.
+  enable_rw_mount "${rootfs}"
+  info "Running e2fsck to check root file system for errors"
+  sudo e2fsck -fn "${rootfs}" ||
+    die "Root file system has errors!"
+  # Flip the bit back so we don't break hashes.
+  disable_rw_mount "${rootfs}"
+}
+
+# Extracts a firmware updater bundle (for firmware image binaries) file
+# (generated by src/platform/firmware/pack_firmware.sh).
+# Args: INPUT_FILE OUTPUT_DIR
+extract_firmware_bundle() {
+  local input="$(readlink -f "$1")"
+  local output_dir="$2"
+  if [ ! -s "${input}" ]; then
+    return 1
+  elif grep -q '^##CUTHERE##' "${input}"; then
+    # Bundle supports self-extraction.
+    "$input" --sb_extract "${output_dir}" ||
+      die "Extracting firmware autoupdate (--sb_extract) failed."
+  else
+    # Legacy bundle - try uudecode.
+    uudecode -o - ${input} | tar -C ${output_dir} -zxf - 2>/dev/null ||
+      die "Extracting firmware autoupdate failed."
+  fi
+}
+
+# Repacks firmware updater bundle content from given folder.
+# Args: INPUT_DIR TARGET_SCRIPT
+repack_firmware_bundle() {
+  local input_dir="$1"
+  local target="$(readlink -f "$2")"
+
+  if [ ! -s "${target}" ]; then
+    return 1
+  elif grep -q '^##CUTHERE##' "${target}"; then
+    # Bundle supports repacking.
+    # Workaround issue crosbug.com/p/33719
+    sed -i \
+      's/shar -Q -q -x -m -w/shar -Q -q -x -m --no-character-count/' \
+      "${target}"
+    "$target" --sb_repack "${input_dir}" ||
+      die "Updating firmware autoupdate (--sb_repack) failed."
+  else
+    # Legacy bundle using uuencode + tar.gz.
+    # Replace MD5 checksum in the firmware update payload.
+    local newfd_checksum="$(md5sum ${input_dir}/bios.bin | cut -f 1 -d ' ')"
+    local temp_version="$(make_temp_file)"
+    cat ${input_dir}/VERSION |
+    sed -e "s#\(.*\)\ \(.*bios.bin.*\)#${newfd_checksum}\ \2#" > ${temp_version}
+    mv ${temp_version} ${input_dir}/VERSION
+
+    # Re-generate firmware_update.tgz and copy over encoded archive in
+    # the original shell ball.
+    sed -ine '/^begin .*firmware_package/,/end/D' "$target"
+    tar zcf - -C "${input_dir}" . |
+      uuencode firmware_package.tgz >>"${target}"
+  fi
+}
+
+# Sign an update payload (usually created by paygen).
+# Args: HASH KEY_DIR OUTPUT
+# Depends on global variable KEY_ORIGIN.
+# <HASH> is supposed to be a SHA256 hash, unencoded, 32 bytes long.
+sign_update_payload() {
+  if [[ "${KEY_ORIGIN}" == "local" ]]; then
+    sign_update_payload_local "$@"
+  elif [[ "${KEY_ORIGIN}" == "kms" ]]; then
+    sign_update_payload_kms "$@"
+  else
+    die "Unsupported KEY_ORIGIN: ${KEY_ORIGIN}"
+  fi
+}
+
+sign_update_payload_local() {
+  local hash=$1
+  local key_dir=$2
+  local output=$3
+  local key_output key_size key_file="${key_dir}/update_key.pem"
+  # Maps key size to verified boot's algorithm id (for pad_digest_utility).
+  # Hashing algorithm is always SHA-256.
+  local algo algos=(
+    [1024]=1
+    [2048]=4
+    [4096]=7
+    [8192]=10
+  )
+
+  key_output=$(futility show "${key_file}")
+  key_size=$(echo "${key_output}" | sed -n '/Key length/s/[^0-9]*//p')
+  algo=${algos[${key_size}]}
+  if [[ -z ${algo} ]]; then
+    die "Unknown algorithm: futility output=${key_output}"
+  fi
+
+  pad_digest_utility ${algo} "${hash}" | \
+    openssl rsautl -sign -pkcs -inkey "${key_file}" -out "${output}"
+}
+
+# Signs a payload with a key stored in Cloud KMS.
+# $key_dir/kms.key must exist and must be a file with content in the
+# following format:
+#   KMS_PROJECT=<project>
+#   KMS_LOCATION=<location>
+#   KMS_KEYRING=<keyring>
+#   KMS_KEY=<key>
+#   KMS_KEYVERSION=<key version>
+sign_update_payload_kms() {
+  local -r hash="$1" key_dir="$2" output="$3"
+  local -r key_file="${key_dir}/kms.key"
+
+  source "${key_file}"
+
+  info "Signing update payload hash ${hash} with key ${key_file}"
+  kms_signer \
+    --project "${KMS_PROJECT}" \
+    --location "${KMS_LOCATION}" \
+    --keyring "${KMS_KEYRING}" \
+    --key "${KMS_KEY}" \
+    --key-version "${KMS_KEYVERSION}" \
+    digest \
+    --input "${hash}" \
+    --output "${output}"
+}
+
+# Sign UEFI binaries, if possible.
+# Args: LOOPDEV
+sign_uefi_binaries() {
+  local loopdev="$1"
+  local kms_option="--nokms"
+  if [[ "${KEY_ORIGIN}" == "kms" ]]; then
+    kms_option="--kms"
+  fi
+
+  local esp_dir
+  if ! esp_dir="$(mount_image_esp "${loopdev}")"; then
+    error "Could not mount EFI partition for signing UEFI binaries"
+    return 1
+  elif [[ -z "${esp_dir}" ]]; then
+    return 0
+  fi
+  "${SCRIPT_DIR}/sign_uefi.sh" -t "${esp_dir}"  -k "${KEY_DIR}" "${kms_option}"
+  sudo umount "${esp_dir}"
+
+  local rootfs_dir="$(make_temp_dir)"
+  mount_loop_image_partition "${loopdev}" 3 "${rootfs_dir}"
+  "${SCRIPT_DIR}/sign_uefi.sh" -t "${rootfs_dir}/boot" -k "${KEY_DIR}" "${kms_option}"
+  sudo umount "${rootfs_dir}"
+
+  info "Signed UEFI binaries"
+  return 0
+}
+
+# Verify the signatures of UEFI binaries.
+# Args: LOOPDEV
+verify_uefi_signatures() {
+  local loopdev="$1"
+  local succeeded=1
+
+  if [[ ! -d "${KEY_DIR}/uefi" ]]; then
+    return 0
+  fi
+
+  local esp_dir
+  if ! esp_dir="$(mount_image_esp "${loopdev}")"; then
+    error "Could not mount EFI partition for verifying UEFI signatures"
+    return 1
+  elif [[ -z "${esp_dir}" ]]; then
+    return 0
+  fi
+  "${SCRIPT_DIR}/verify_uefi.sh" "${esp_dir}" "${esp_dir}" \
+      "${KEY_DIR}/uefi" || succeeded=0
+
+  local rootfs_dir="$(make_temp_dir)"
+  mount_loop_image_partition_ro "${loopdev}" 3 "${rootfs_dir}"
+  "${SCRIPT_DIR}/verify_uefi.sh" "${rootfs_dir}/boot" "${esp_dir}" \
+      "${KEY_DIR}/uefi" || succeeded=0
+  sudo umount "${rootfs_dir}"
+
+  sudo umount "${esp_dir}"
+
+  if [[ "${succeeded}" == "0" ]]; then
+    die "UEFI signature verification failed"
+  fi
+}
+
+# Verify an image including rootfs hash using the specified keys.
+verify_image() {
+  local loopdev=$(loopback_partscan "${INPUT_IMAGE}")
+  local loop_rootfs="${loopdev}p3"
+
+  info "Verifying RootFS hash..."
+  # What we get from image.
+  local kernel_config
+  # What we calculate from the rootfs.
+  local new_kernel_config
+  # Depending on the type of image, the verity parameters may
+  # exist in either kernel partition 2 or kernel partition 4
+  local partnum
+  for partnum in 2 4; do
+    info "Considering Kernel partition ${partnum}"
+    kernel_config=$(sudo dump_kernel_config "${loopdev}p${partnum}")
+    local hash_image=$(make_temp_file)
+    if ! calculate_rootfs_hash "${loop_rootfs}" "${kernel_config}" \
+      "${hash_image}"; then
+      info "Trying next kernel partition."
+      continue
+    fi
+    new_kernel_config="$CALCULATED_KERNEL_CONFIG"
+    break
+  done
+
+  # Note: If calculate_rootfs_hash succeeded above, these should
+  # be non-empty.
+  expected_hash=$(get_hash_from_config "${new_kernel_config}")
+  got_hash=$(get_hash_from_config "${kernel_config}")
+
+  if [ -z "${expected_hash}" ] || [ -z "${got_hash}" ]; then
+    die "Couldn't verify RootFS hash on the image."
+  fi
+
+  if [ ! "${got_hash}" = "${expected_hash}" ]; then
+    cat <<EOF
+FAILED: RootFS hash is incorrect.
+Expected: ${expected_hash}
+Got: ${got_hash}
+EOF
+    exit 1
+  else
+    info "PASS: RootFS hash is correct (${expected_hash})"
+  fi
+
+  # Now try and verify kernel partition signature.
+  set +e
+  local try_key=${KEY_DIR}/recovery_key.vbpubk
+  info "Testing key verification..."
+  # The recovery key is only used in the recovery mode.
+  echo -n "With Recovery Key (Recovery Mode ON, Dev Mode OFF): " && \
+  { load_kernel_test "${INPUT_IMAGE}" "${try_key}" -b 2 >/dev/null 2>&1 && \
+    echo "YES"; } || echo "NO"
+  echo -n "With Recovery Key (Recovery Mode ON, Dev Mode ON): " && \
+  { load_kernel_test "${INPUT_IMAGE}" "${try_key}" -b 3 >/dev/null 2>&1 && \
+    echo "YES"; } || echo "NO"
+
+  try_key=${KEY_DIR}/kernel_subkey.vbpubk
+  # The SSD key is only used in non-recovery mode.
+  echo -n "With SSD Key (Recovery Mode OFF, Dev Mode OFF): " && \
+  { load_kernel_test "${INPUT_IMAGE}" "${try_key}" -b 0 >/dev/null 2>&1  && \
+    echo "YES"; } || echo "NO"
+  echo -n "With SSD Key (Recovery Mode OFF, Dev Mode ON): " && \
+  { load_kernel_test "${INPUT_IMAGE}" "${try_key}" -b 1 >/dev/null 2>&1 && \
+    echo "YES"; } || echo "NO"
+  set -e
+
+  verify_image_rootfs "${loop_rootfs}"
+
+  verify_uefi_signatures "${INPUT_IMAGE}"
+}
+
+# Update the legacy bootloader templates in EFI partition if available.
+# Args: LOOPDEV KERNEL
+update_legacy_bootloader() {
+  local loopdev="$1"
+  local loop_kern="$2"
+
+  local esp_dir
+  if ! esp_dir="$(mount_image_esp "${loopdev}")"; then
+    error "Could not mount EFI partition for updating legacy bootloader cfg."
+    return 1
+  elif [[ -z "${esp_dir}" ]]; then
+    info "Not updating legacy bootloader configs: ${loopdev}"
+    return 0
+  fi
+
+  # If we can't find the dm parameter in the kernel config, bail out now.
+  local kernel_config=$(sudo dump_kernel_config "${loop_kern}")
+  local root_hexdigest="$(get_hash_from_config "${kernel_config}")"
+  if [[ -z "${root_hexdigest}" ]]; then
+    error "Couldn't grab root_digest from kernel partition ${loop_kern}"
+    error " (config: ${kernel_config})"
+    return 1
+  fi
+  # Update syslinux configs for legacy BIOS systems.
+  if [[ -d "${esp_dir}/syslinux" ]]; then
+    local cfg=("${esp_dir}"/syslinux/*.cfg)
+    if ! sudo sed -i -r \
+      "s/\broot_hexdigest=[a-z0-9]+/root_hexdigest=${root_hexdigest}/g" \
+      "${cfg[@]}"; then
+        error "Updating syslinux configs failed: '${cfg[*]}'"
+        return 1
+    fi
+  fi
+  # Update grub configs for EFI systems.
+  local grub_cfg="${esp_dir}/efi/boot/grub.cfg"
+  if [[ -f "${grub_cfg}" ]]; then
+    if ! sudo sed -i -r \
+      "s/\broot_hexdigest=[a-z0-9]+/root_hexdigest=${root_hexdigest}/g" \
+      "${grub_cfg}"; then
+        error "Updating grub config failed: '${grub_cfg}'"
+        return 1
+    fi
+  fi
+}
+
+# Sign an image file with proper keys.
+# Args: IMAGE_TYPE INPUT OUTPUT DM_PARTNO KERN_A_KEYBLOCK KERN_A_PRIVKEY \
+#       KERN_B_KEYBLOCK KERN_B_PRIVKEY
+#
+# A ChromiumOS image file (INPUT) always contains 2 partitions (kernel A & B).
+# This function will rebuild hash data by DM_PARTNO, resign kernel partitions by
+# their KEYBLOCK and PRIVKEY files, and then write to OUTPUT file. Note some
+# special images (specified by IMAGE_TYPE, like 'recovery' or 'factory_install')
+# may have additional steps (ex, tweaking verity hash or not stripping files)
+# when generating output file.
+sign_image_file() {
+  local image_type="$1"
+  local input="$2"
+  local output="$3"
+  local dm_partno="$4"
+  local kernA_keyblock="$5"
+  local kernA_privkey="$6"
+  local kernB_keyblock="$7"
+  local kernB_privkey="$8"
+
+  info "Preparing ${image_type} image..."
+  cp --sparse=always "${input}" "${output}"
+
+  local loopdev=$(loopback_partscan "${output}")
+  local loop_kern="${loopdev}p${dm_partno}"
+  local loop_rootfs="${loopdev}p3"
+
+  sign_uefi_binaries "${loopdev}"
+  # We do NOT strip /boot for factory installer, since some devices need it to
+  # boot EFI. crbug.com/260512 would obsolete this requirement.
+  #
+  # We also do NOT strip /boot for legacy BIOS or EFI devices.  This is because
+  # "cros_installer postinst" on BIOS or EFI systems relies on presence of
+  # /boot in rootfs to update kernel.  We infer the BIOS type from the kernel
+  # config.
+  local loop_kerna="${loopdev}p2"
+  local kerna_config="$(sudo dump_kernel_config "${loop_kerna}")"
+  if [[ "${image_type}" != "factory_install" &&
+        " ${kerna_config} " != *" cros_legacy "* &&
+        " ${kerna_config} " != *" cros_efi "* ]]; then
+    "${SCRIPT_DIR}/strip_boot_from_image.sh" --image "${loop_rootfs}"
+  fi
+  update_rootfs_hash "${loopdev}" "${loop_kern}" \
+    "${kernA_keyblock}" "${kernA_privkey}" \
+    "${kernB_keyblock}" "${kernB_privkey}"
+  update_stateful_partition_vblock "${loopdev}"
+  if ! update_legacy_bootloader "${loopdev}" "${loop_kern}"; then
+    # Error is already logged.
+    return 1
+  fi
+  # Let non-root users read it. It's okay.
+  sudo chmod 0644 "${output}"
+  info "Signed ${image_type} image output to ${output}"
+}
+
+# Verification
+case ${TYPE} in
+dump_config)
+  check_argc $# 2
+  loopdev=$(loopback_partscan "${INPUT_IMAGE}")
+  for partnum in 2 4; do
+    info "kernel config in partition number ${partnum}:"
+    sudo dump_kernel_config "${loopdev}p${partnum}"
+    echo
+  done
+  exit 0
+  ;;
+verify)
+  check_argc $# 2
+  verify_image
+  exit 0
+  ;;
+*)
+  # All other signing commands take 4 to 5 args.
+  if [ -z "${OUTPUT_IMAGE}" ]; then
+    # Friendlier message.
+    usage "Missing output image name"
+  fi
+  check_argc $# 4 5
+  ;;
+esac
+
+# If a version file was specified, read the firmware and kernel
+# versions from there.
+if [ -n "${VERSION_FILE}" ]; then
+  FIRMWARE_VERSION=$(sed -n 's#^firmware_version=\(.*\)#\1#pg' ${VERSION_FILE})
+  KERNEL_VERSION=$(sed -n 's#^kernel_version=\(.*\)#\1#pg' ${VERSION_FILE})
+fi
+info "Using firmware version: ${FIRMWARE_VERSION}"
+info "Using kernel version: ${KERNEL_VERSION}"
+
+# Make all modifications on output copy.
+if [[ "${TYPE}" == "base" ]]; then
+  sign_image_file "SSD" "${INPUT_IMAGE}" "${OUTPUT_IMAGE}" 2 \
+    "${KEY_DIR}/kernel.keyblock" "${KEY_DIR}/kernel_data_key.vbprivk" \
+    "${KEY_DIR}/kernel.keyblock" "${KEY_DIR}/kernel_data_key.vbprivk"
+elif [[ "${TYPE}" == "update_payload" ]]; then
+  # The argument names here are a little awkard because sign_update_payload
+  # doesn't sign "image" but only signs hashes, but we want to use the same
+  # interface as sign_image_file, so ...
+  sign_update_payload ${INPUT_IMAGE} ${KEY_DIR} ${OUTPUT_IMAGE}
+else
+  die "Invalid type ${TYPE}"
+fi