#!/bin/bash
#
# Copyright 2024 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

SCRIPT_BASE="$(dirname "$0")"
. "$SCRIPT_BASE/common_minimal.sh"
load_shflags || exit 1

FLAGS_HELP="
Swap the EC RW (ecrw) within an AP firmware (BIOS) image.
"

# Flags.
DEFINE_string image "" "The AP firmware file (e.g 'image-steelix.bin') to swap out ecrw" i
DEFINE_string ec "" "The EC firmware file (e.g 'ec.bin')" e

# Parse command line.
FLAGS "$@" || exit 1
eval set -- "${FLAGS_ARGV}"

# Only after this point should you enable `set -e` as shflags does not work
# when that is turned on first.
set -e

FMAP_REGIONS=( "FW_MAIN_A" "FW_MAIN_B" )
CBFS_ECRW_NAME="ecrw"
CBFS_ECRW_HASH_NAME="ecrw.hash"
CBFS_ECRW_VERSION_NAME="ecrw.version"
CBFS_ECRW_CONFIG_NAME="ecrw.config"

swap_ecrw() {
  local ap_file=$1
  local ec_file=$2
  local temp_dir
  local info
  local ecrw_file
  local ecrw_hash_file
  local ecrw_ver_file
  local ecrw_comp_type
  local ecrw_ver
  local apro_ver
  local aprw_ver
  temp_dir=$(mktemp -d)
  ecrw_file="${temp_dir}/ecrw"
  futility dump_fmap -x "${ec_file}" "RW_FW:${ecrw_file}" >/dev/null
  info "EC RW extracted to ${ecrw_file}"

  ecrw_hash_file="${temp_dir}/ecrw.hash"
  openssl dgst -sha256 -binary "${ecrw_file}" > "${ecrw_hash_file}"
  info "EC RW hash saved to ${ecrw_hash_file}"

  ecrw_ver_file="${temp_dir}/ecrw.version"
  futility dump_fmap -x "${ec_file}" "RW_FWID:${ecrw_ver_file}" >/dev/null

  for region in "${FMAP_REGIONS[@]}"
  do
    info="$(cbfstool "${ap_file}" print -r "${region}" -k -v \
      | grep -m 1 "^${CBFS_ECRW_NAME}\s")"
    ecrw_comp_type="$(cut -f7- <<< "${info}" | grep -o '\<comp\>:\w*' \
      | cut -d: -f2)"
    ecrw_comp_type=${ecrw_comp_type:-none}
    cbfstool "${ap_file}" remove -r "${region}" -n "${CBFS_ECRW_NAME}"
    cbfstool "${ap_file}" remove -r "${region}" -n "${CBFS_ECRW_HASH_NAME}"
    cbfstool "${ap_file}" remove -r "${region}" -n "${CBFS_ECRW_VERSION_NAME}" \
      || warn "${CBFS_ECRW_VERSION_NAME} not found, but will be added"
    # TODO(b/307788351): Update ecrw.config. Right now the config info cannot
    # be obtained from ec.bin.
    cbfstool "${ap_file}" remove -r "${region}" \
      -n "${CBFS_ECRW_CONFIG_NAME_NAME}" || true
    cbfstool "${ap_file}" expand -r "${region}"
    cbfstool "${ap_file}" add -r "${region}" -t raw \
      -c "${ecrw_comp_type}" -f "${ecrw_file}" -n "${CBFS_ECRW_NAME}"
    cbfstool "${ap_file}" add -r "${region}" -t raw \
      -c none -f "${ecrw_hash_file}" -n "${CBFS_ECRW_HASH_NAME}"
    cbfstool "${ap_file}" add -r "${region}" -t raw \
      -c none -f "${ecrw_ver_file}" -n "${CBFS_ECRW_VERSION_NAME}"
  done

  local keyset
  for keyset in /usr/share/vboot/devkeys "${SCRIPT_BASE}/../../tests/devkeys"; do
    [[ -d "${keyset}" ]] && break
  done

  # 'futility sign' will call 'cbfstool truncate' if needed
  futility sign "${ap_file}" --keyset "${keyset}"

  ecrw_ver=$(futility update --manifest -e "${ec_file}" \
    | jq -r '.default.ec.versions.rw')
  apro_ver=$(futility update --manifest -i "${ap_file}" \
    | jq -r '.default.host.versions.ro')
  aprw_ver=$(futility update --manifest -i "${ap_file}" \
    | jq -r '.default.host.versions.rw')
  info "${CBFS_ECRW_NAME} (${ecrw_ver}) swapped in ${ap_file} (RO:${apro_ver}, RW:${aprw_ver})"
  info "Done"
}

main() {
  if [[ -z "${FLAGS_image}" ]]; then
    flags_help
    die "-i or --image required."
  fi
  if [[ -z "${FLAGS_ec}" ]]; then
    flags_help
    die "-e or --ec required."
  fi

  swap_ecrw "${FLAGS_image}" "${FLAGS_ec}"
}

main "$@"