#!/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
SERVICE_ACCOUNT=$6
VERSION_FILE=$7

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>
# dm-verity old version(0) has key value pair for all verity args
# Ex: "dm=1 vroot none ro 1,0 4077568 verity
# payload=PARTUUID=62878627-9962-574A-9B44-1A231DBFD5B5
# hashtree=PARTUUID=62878627-9962-574A-9B44-1A231DBFD5B5 hashstart=4077568
# alg=sha256
# root_hexdigest=85796f88c1d2eb6c108881054a234afc1b6dad2851324c6f5b1be81aa0ffa14
# salt=d5582575b3993bbb0d30bd16bd74e0011ccbbbcee58e4757dce2d8ed6c5226e8"
#
# dm-verity new version(1) has all arguments as positional arguments.
# Ex: "dm-mod.create=vroot,,,ro,0 4077568 verity 0
# PARTUUID=2541852D-1C88-9443-BDB3-81E212A0CED1
# PARTUUID=2541852D-1C88-9443-BDB3-81E212A0CED1 4096 4096 509696 509696 sha256
# e68c29018b3b6923282990fbd5ea2c35fb5e4c22c9efe944eb84f2dfb8732daf
# 7389b4ebde0f5a43c4ea0e54adbc1258765651596df5beee6b4503152aa64684"
get_verity_arg() {
  local verity_arg=$1
  local arg=$2
  if [[ ${dm_verity_version} -eq 0 ]]
  then
    echo "$1" | sed -n "s/.*\b$2=\([^ \"]*\).*/\1/p"
  elif [[ ${dm_verity_version} -eq 1 ]]
  then
    case "$arg" in
      payload)
        echo "${verity_arg}" | awk '{print $5}'
        ;;
      hashtree)
        echo "${verity_arg}" | awk '{print $6}'
        ;;
      hashstart)
        hashstart=$(echo "${verity_arg}" | awk '{print $10}')
        hashstart=$(($hashstart<<3))
        echo "${hashstart}"
        ;;
      alg)
        echo "${verity_arg}" | awk '{print $11}'
        ;;
      root_hexdigest)
        echo "${verity_arg}" | awk '{print $12}'
        ;;
      salt)
        echo "${verity_arg}" | awk '{print $13}'
        ;;
    esac
  fi
}

# Get the dmparams parameters from a kernel config.
get_dmparams_from_config() {
  local kernel_cmd_line=$1
  echo ${kernel_cmd_line} | sed -nre "s/.*$dm_str\"([^\"]*)\".*/\1/p"
}
# Get the verity root digest hash from a kernel config command line.
get_hash_from_config() {
  local kernel_cmd_line=$1
  local dm_config=$(get_dmparams_from_config "${kernel_cmd_line}")
  local vroot_dev=$(get_dm_slave "${dm_config}" vroot)
  echo $(get_verity_arg "${vroot_dev}" root_hexdigest)
}

get_alg_from_config() {
  local kernel_cmd_line=$1
  local dm_config=$(get_dmparams_from_config "${kernel_cmd_line}")
  local vroot_dev=$(get_dm_slave "${dm_config}" vroot)
  echo $(get_verity_arg "${vroot_dev}" alg)
}

# 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
  if [[ $dm_verity_version -eq 0 ]]
  then
    echo $(echo "${dm}" | sed -nre "s/.*${device}[^,]*,([^,]*).*/\1/p")
  elif [[ $dm_verity_version -eq 1 ]]
  then
    echo $(echo "${dm}" | awk -F, '{print $NF}')
  fi
}

# 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_CMD_LINE=
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_cmd_line=$2
  local hash_image=$3
  local dm_config=$(get_dmparams_from_config "${kernel_cmd_line}")

  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} version="${dm_verity_version}")
  # 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}|")"
  if [[ ${dm_verity_version} -eq 0 ]]
  then
    CALCULATED_DM_ARGS="$(set_dm_slave "${dm_config}" vroot "${slave}")"
    CALCULATED_KERNEL_CMD_LINE="$(echo "${kernel_cmd_line}" |
    sed -e 's#\(.*dm="\)\([^"]*\)\(".*\)'"#\1${CALCULATED_DM_ARGS}\3#g")"
  elif [[ ${dm_verity_version} -eq 1 ]]
  then
    CALCULATED_DM_ARGS="vroot,,,ro,""${slave}"
    CALCULATED_KERNEL_CMD_LINE="$(echo "${kernel_cmd_line}" |
    sed -e "s#\(.*$dm_str\"\)\([^\"]*\)\(\".*\)""#\1${CALCULATED_DM_ARGS}\3#g")"
  fi
}

# 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_cmd_line=$(sudo dump_kernel_config "${loop_kern}")
  local dm_config=$(get_dmparams_from_config "${kernel_cmd_line}")
  if [ -z "${dm_config}" ]; then
    error "Couldn't grab dm_config from kernel ${loop_kern}"
    error " (config: ${kernel_cmd_line})"
    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_cmd_line}" \
    "${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_cmd_line=

  for kernelpart in 2 4; do
    loop_kern="${loopdev}p${kernelpart}"
    if ! new_kernel_cmd_line="$(
         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_cmd_line="$(echo "${new_kernel_cmd_line}" |
      sed -e "s#\(.*$dm_str\"\)\([^\"]*\)\(\".*\)""#\1${dm_args}\3#g")"
    info "New config for kernel partition ${kernelpart} is:"
    echo "${new_kernel_cmd_line}" | 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" service_account="$4"
  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}" \
    --service-account "${service_account}"
}

# Sign UEFI binaries, if possible.
# Args: LOOPDEV
sign_uefi_binaries() {
  local loopdev="$1"
  local kms_option="--nokms"
  local service_account="$2"
  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}" --service_account "${service_account}"
  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}" --service_account "${service_account}"
  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_cmd_line
  # What we calculate from the rootfs.
  local new_kernel_cmd_line
  # 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_cmd_line=$(sudo dump_kernel_config "${loopdev}p${partnum}")
    local hash_image=$(make_temp_file)
    if ! calculate_rootfs_hash "${loop_rootfs}" "${kernel_cmd_line}" \
      "${hash_image}"; then
      info "Trying next kernel partition."
      continue
    fi
    new_kernel_cmd_line="$CALCULATED_KERNEL_CMD_LINE"
    break
  done

  # Note: If calculate_rootfs_hash succeeded above, these should
  # be non-empty.
  expected_hash=$(get_hash_from_config "${new_kernel_cmd_line}")
  got_hash=$(get_hash_from_config "${kernel_cmd_line}")

  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_cmd_line=$(sudo dump_kernel_config "${loop_kern}")
  local root_hexdigest="$(get_hash_from_config "${kernel_cmd_line}")"
  local alg="$(get_alg_from_config "${kernel_cmd_line}")"
  if [[ -z "${root_hexdigest}" ]]; then
    error "Couldn't grab root_digest from kernel partition ${loop_kern}"
    error " (config: ${kernel_cmd_line})"
    return 1
  fi
  # Update syslinux configs for legacy BIOS systems.
  if [[ -d "${esp_dir}/syslinux" ]]; then
    local cfg=("${esp_dir}"/syslinux/*.cfg)
    if [[ $dm_verity_version -eq 0 ]]
    then
      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
    elif [[ $dm_verity_version -eq 1 ]]
    then
      if ! sudo sed -i -r \
        "s/${alg} [a-f0-9]+/${alg} ${root_hexdigest}/g" \
        "${cfg[@]}"; then
          error "Updating syslinux configs failed: '${cfg[*]}'"
          return 1
      fi
    fi
  fi
  # Update grub configs for EFI systems.
  local grub_cfg="${esp_dir}/efi/boot/grub.cfg"
  if [[ -f "${grub_cfg}" ]]; then
    if [[ $dm_verity_version -eq 0 ]]
    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
    elif [[ $dm_verity_version -eq 1 ]]
    then
      if ! sudo sed -i -r \
        "s/${alg} [a-f0-9]+/${alg} ${root_hexdigest}/g" \
        "${grub_cfg}"; then
          error "Updating grub config failed: '${grub_cfg}'"
          return 1
      fi
    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"
  local service_account="$9"

  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}" "${service_account}"
  # 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 echo "${kerna_config}" | grep -q "dm="
  then
    dm_str="dm="
    dm_verity_version=0
  else
    dm_str="dm-mod.create="
    dm_verity_version=1
  fi
  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 6 args.
  if [ -z "${OUTPUT_IMAGE}" ]; then
    # Friendlier message.
    usage "Missing output image name"
  fi
  check_argc $# 4 6
  ;;
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" ${SERVICE_ACCOUNT}
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} ${SERVICE_ACCOUNT}
else
  die "Invalid type ${TYPE}"
fi