sign_uefi.sh:Port google3 changes required for COS

As part of the Cusky signing design, we carry custom changes to the
sign_uefi.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.

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

BUG=b/170128847
TEST=presubmit; merged signer container with BE container, patched these
changes in and ran this script (via sign_official_build.sh in
go/cos-rev/c/18671)

Change-Id: Id3ecd8b5cf056d244f8cfd89f65d9f13322b0998
Reviewed-on: https://cos-review.googlesource.com/c/third_party/platform/vboot_reference/+/18670
Tested-by: Cusky Presubmit Bot <presubmit@cos-infra-prod.iam.gserviceaccount.com>
Reviewed-by: Robert Kolchmeyer <rkolchmeyer@google.com>
Reviewed-on: https://cos-review.googlesource.com/c/third_party/platform/vboot_reference/+/21902
Tested-by: Cusky QA Presubmit Bot <presubmit@cos-infra-prototype.iam.gserviceaccount.com>
diff --git a/scripts/image_signing/sign_uefi.sh b/scripts/image_signing/sign_uefi.sh
index 6deb280..14c328e 100755
--- a/scripts/image_signing/sign_uefi.sh
+++ b/scripts/image_signing/sign_uefi.sh
@@ -4,34 +4,64 @@
 # found in the LICENSE file.
 
 . "$(dirname "$0")/common.sh"
+load_shflags || exit 1
+
+DEFINE_string target_dir "" "Directory to put signed file in" "t"
+DEFINE_string key_dir "" "Directory of signing keys and certificates" "k"
+DEFINE_boolean kms $FLAGS_FALSE "Whether or not to sign with KMS keys" ""
+
+FLAGS "$@" || exit 1
+eval set -- "$FLAGS_ARGV"
 
 set -e
 
-usage() {
-  cat <<EOF
-Usage: $PROG /path/to/target/dir /path/to/uefi/keys/dir
+# Resigns a signed EFI file with key in Cloud KMS.
+# Requires a tool called `kms_signer` on the system path.
+# The second argument, kms_key, should point to a file with content in the
+# following format:
+#   KMS_PROJECT=<project>
+#   KMS_LOCATION=<location>
+#   KMS_KEYRING=<keyring>
+#   KMS_KEY=<key>
+#   KMS_KEYVERSION=<key version>
+resign_with_kms() {
+  local -r target="$1" kms_key="$2" kms_cert="$3" kms_ca_cert="$4"
+  local old_sig="$(mktemp)" new_sig="$(mktemp)" resigned="$(mktemp)"
 
-Sign the UEFI binaries in the target directory.
-The target directory can be either the root of ESP or /boot of root filesystem.
-EOF
-  if [[ $# -gt 0 ]]; then
-    error "$*"
-    exit 1
-  fi
-  exit 0
+  source "${kms_key}"
+
+  # Detach the signature and resign.
+  info "Resigning EFI file ${target} with key ${kms_key} and certificate ${kms_cert}"
+  sbattach --detach "${old_sig}" "${target}"
+  kms_signer \
+    --project "${KMS_PROJECT}" \
+    --location "${KMS_LOCATION}" \
+    --keyring "${KMS_KEYRING}" \
+    --key "${KMS_KEY}" \
+    --key-version "${KMS_KEYVERSION}" \
+    pkcs7 \
+    --signing-cert "${kms_cert}" \
+    --input "${old_sig}" \
+    --output "${new_sig}"
+
+  cp "${target}" "${resigned}"
+  sbattach --attach "${new_sig}" "${resigned}"
+  mv "${resigned}" "${target}"
+  sbverify --cert "${kms_ca_cert}" "${target}"
+  rm -f "${old_sig}" "${new_sig}"
 }
 
 # Signs an EFI binary file, if possible.
-# Args: TARGET_FILE TEMP_DIR PRIVATE_KEY SIGN_CERT VERIFY_CERT
+# Args: TARGET_FILE TEMP_DIR PRIVATE_KEY SIGN_CERT VERIFY_CERT [KMS_KEY] [KMS_CERT] [KMS_CA_CERT]
 sign_efi_file() {
   local target="$1"
   local temp_dir="$2"
   local priv_key="$3"
   local sign_cert="$4"
   local verify_cert="$5"
-  if [[ -z "${verify_cert}" ]]; then
-    verify_cert="${sign_cert}"
-  fi
+  local kms_key="$6"
+  local kms_cert="$7"
+  local kms_ca_cert="$8"
 
   info "Signing efi file ${target}"
   sudo sbattach --remove "${target}" || true
@@ -42,45 +72,61 @@
     sudo cp -f "${signed_file}" "${target}"
     sbverify --cert "${verify_cert}" "${target}" || die "Verification failed"
   fi
+
+  if [[ "${FLAGS_kms}" == "${FLAGS_TRUE}" ]]; then
+    resign_with_kms "${target}" "${kms_key}" "${kms_cert}" "${kms_ca_cert}"
+  fi
 }
 
 main() {
-  local target_dir="$1"
-  local key_dir="$2"
+  local kms_key kms_cert kms_ca_cert
 
-  if [[ $# -ne 2 ]]; then
-    usage "command takes exactly 2 args"
+  local prereqs=(sbattach sbsign sbverify)
+  if [[ "${FLAGS_kms}" == "${FLAGS_TRUE}" ]]; then
+    prereqs+=(kms_signer)
   fi
+  for prereq in ${prereqs[@]}; do
+  if ! type -P "${prereq}" &>/dev/null; then
+    die "Prerequisite not found: ${prereq}."
+  fi
+  done
 
-  if ! type -P sbattach &>/dev/null; then
-    die "Cannot sign UEFI binaries (sbattach not found)."
-  fi
-  if ! type -P sbsign &>/dev/null; then
-    die "Cannot sign UEFI binaries (sbsign not found)."
-  fi
-  if ! type -P sbverify &>/dev/null; then
-    die "Cannot sign UEFI binaries (sbverify not found)."
-  fi
+  local bootloader_dir="${FLAGS_target_dir}/efi/boot"
+  local syslinux_dir="${FLAGS_target_dir}/syslinux"
+  local kernel_dir="${FLAGS_target_dir}"
 
-  local bootloader_dir="${target_dir}/efi/boot"
-  local syslinux_dir="${target_dir}/syslinux"
-  local kernel_dir="${target_dir}"
-
-  local verify_cert="${key_dir}/db/db.pem"
+  local verify_cert="${FLAGS_key_dir}/db/db.pem"
   if [[ ! -f "${verify_cert}" ]]; then
     die "No verification cert: ${verify_cert}"
   fi
 
-  local sign_cert="${key_dir}/db/db.children/db_child.pem"
+  local sign_cert="${FLAGS_key_dir}/db/db.children/db_child.pem"
   if [[ ! -f "${sign_cert}" ]]; then
     die "No signing cert: ${sign_cert}"
   fi
 
-  local sign_key="${key_dir}/db/db.children/db_child.rsa"
+  local sign_key="${FLAGS_key_dir}/db/db.children/db_child.rsa"
   if [[ ! -f "${sign_key}" ]]; then
     die "No signing key: ${sign_key}"
   fi
 
+  if [[ "${FLAGS_kms}" == "${FLAGS_TRUE}" ]]; then
+    kms_key="${FLAGS_key_dir}/kms/db_child.key"
+    if [[ ! -f "${kms_key}" ]]; then
+      die "No KMS key: ${kms_key}"
+    fi
+
+    kms_cert="${FLAGS_key_dir}/kms/db_child.crt"
+    if [[ ! -f "${kms_cert}" ]]; then
+      die "No KMS cert: ${kms_cert}"
+    fi
+
+    kms_ca_cert="${FLAGS_key_dir}/kms/db.crt"
+    if [[ ! -f "${kms_ca_cert}" ]]; then
+      die "No KMS CA cert: ${kms_ca_cert}"
+    fi
+  fi
+
   local working_dir="$(make_temp_dir)"
 
   local efi_file
@@ -89,7 +135,8 @@
       continue
     fi
     sign_efi_file "${efi_file}" "${working_dir}" \
-        "${sign_key}" "${sign_cert}" "${verify_cert}"
+        "${sign_key}" "${sign_cert}" "${verify_cert}" \
+        "${kms_key}" "${kms_cert}" "${kms_ca_cert}"
   done
 
   local syslinux_kernel_file
@@ -98,13 +145,15 @@
       continue
     fi
     sign_efi_file "${syslinux_kernel_file}" "${working_dir}" \
-        "${sign_key}" "${sign_cert}" "${verify_cert}"
+        "${sign_key}" "${sign_cert}" "${verify_cert}" \
+        "${kms_key}" "${kms_cert}" "${kms_ca_cert}"
   done
 
   local kernel_file="$(readlink -f "${kernel_dir}/vmlinuz")"
   if [[ -f "${kernel_file}" ]]; then
     sign_efi_file "${kernel_file}" "${working_dir}" \
-        "${sign_key}" "${sign_cert}" "${verify_cert}"
+        "${sign_key}" "${sign_cert}" "${verify_cert}" \
+        "${kms_key}" "${kms_cert}" "${kms_ca_cert}"
   fi
 }