#!/bin/bash
# Copyright 2018 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.

. "$(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

# 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)"

  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 [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"
  local kms_key="$6"
  local kms_cert="$7"
  local kms_ca_cert="$8"

  info "Signing efi file ${target}"
  sudo sbattach --remove "${target}" || true
  local signed_file="${temp_dir}/$(basename "${target}")"
  sbsign --key="${priv_key}" --cert="${sign_cert}" \
      --output="${signed_file}" "${target}" || warn "Cannot sign ${target}"
  if [[ -f "${signed_file}" ]]; then
    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 kms_key kms_cert kms_ca_cert

  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

  local bootloader_dir="${FLAGS_target_dir}/efi/boot"
  local syslinux_dir="${FLAGS_target_dir}/syslinux"
  local kernel_dir="${FLAGS_target_dir}"

  local verify_cert="${FLAGS_key_dir}/db/db.pem"
  if [[ ! -f "${verify_cert}" ]]; then
    die "No verification cert: ${verify_cert}"
  fi

  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="${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
  for efi_file in "${bootloader_dir}"/*.efi; do
    if [[ ! -f "${efi_file}" ]]; then
      continue
    fi
    sign_efi_file "${efi_file}" "${working_dir}" \
        "${sign_key}" "${sign_cert}" "${verify_cert}" \
        "${kms_key}" "${kms_cert}" "${kms_ca_cert}"
  done

  local syslinux_kernel_file
  for syslinux_kernel_file in "${syslinux_dir}"/vmlinuz.?; do
    if [[ ! -f "${syslinux_kernel_file}" ]]; then
      continue
    fi
    sign_efi_file "${syslinux_kernel_file}" "${working_dir}" \
        "${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}" \
        "${kms_key}" "${kms_cert}" "${kms_ca_cert}"
  fi
}

main "$@"