#!/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
DEFINE_string ec_config "" "The EC config file (default is 'ec.config')"
DEFINE_string ap_for_ec "" "The AP firmware file (e.g 'image-steelix.bin') as source of EC firmware file" a
DEFINE_string raw_ecrw "" "The raw EC RW file (e.g. 'ec.RW.flat')" r
DEFINE_string ec_ver "" "The EC version file (e.g. 'ec.version')" v

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

cbfstool_check_exist() {
  local ap_file="$1"
  local region="$2"
  local name="$3"
  cbfstool "${ap_file}" print -r "${region}" -k | grep -q "^${name}"$'\t'
}

cbfstool_try_extract() {
  local ap_file="$1"
  local region="$2"
  local name="$3"
  local output="$4"

  if cbfstool_check_exist "${ap_file}" "${region}" "${name}"; then
    if [[ -e "${output}" ]]; then
      die "Extracting should not override file. ${output} already exists."
    fi
    cbfstool "${ap_file}" extract -r "${region}" -n "${name}" -f "${output}"
    return 0
  fi
  return 1
}

cbfstool_try_remove() {
  local ap_file="$1"
  local region="$2"
  local name="$3"

  if cbfstool_check_exist "${ap_file}" "${region}" "${name}"; then
    cbfstool "${ap_file}" remove -r "${region}" -n "${name}"
    return 0
  fi
  return 1
}

extract_ecrw_files_from_ap() {
  local ecrw_file=$1
  local ecrw_ver_file=$2
  local ecrw_config_file=$3
  local ap_for_ec_file=$4

  local region="${FMAP_REGIONS[0]}"

  cbfstool "${ap_for_ec_file}" extract -r "${region}" -n "${CBFS_ECRW_NAME}" \
    -f "${ecrw_file}"
  info "EC RW extracted to ${ecrw_file}"

  cbfstool_try_extract "${ap_for_ec_file}" "${region}" \
    "${CBFS_ECRW_VERSION_NAME}" "${ecrw_ver_file}" \
    || warn "${CBFS_ECRW_VERSION_NAME} not found in source AP file."

  cbfstool_try_extract "${ap_for_ec_file}" "${region}" \
    "${CBFS_ECRW_CONFIG_NAME}" "${ecrw_config_file}" \
    || warn "${CBFS_ECRW_CONFIG_NAME} not found in source AP file."
}

# Dump the 32-bit hex string from the given offset in the given EC-RW file.
read_hex32() {
  local ecrw_file=$1
  local offset=$2

  hexdump -n 4 -s "${offset}" -e '1/4 "%08x"' "${ecrw_file}" 2>/dev/null
}

truncate_ec_rw() {
  local ecrw_file=$1
  local offset=0
  local max_offset=4096

  # The image_data is defined in src/platform/ec/include/cros_version.h:
  #   struct image_data {
  #     uint32_t cookie1;
  #     char version[32];
  #     uint32_t size;
  #     int32_t rollback_version;
  #     uint32_t cookie2;
  #     ...
  #   } __packed;
  readonly COOKIE1_SIGNATURE="ce778899"
  readonly COOKIE2_SIGNATURE="ceaabbdd"
  readonly COOKIE2_OFFSET=44
  readonly SIZE_OFFSET=36

  # Search for the cookies of the image_data header
  while true; do
    if [[ "${offset}" -gt "${max_offset}" ]]; then
      die "No image_data's cookies found within ${max_offset} bytes."
    fi

    cookie1="$(read_hex32 "${ecrw_file}" "${offset}")"
    if [[ "${cookie1}" == "${COOKIE1_SIGNATURE}" ]]; then
      local cookie2_offset=$((offset + COOKIE2_OFFSET))
      cookie2="$(read_hex32 "${ecrw_file}" "${cookie2_offset}")"
      if [[ "${cookie2}" == "${COOKIE2_SIGNATURE}" ]]; then
        break
      fi
    fi
    offset=$((offset + 4))
  done

  # Retrieve image_data.size
  local size_hex
  size_hex="$(read_hex32 "${ecrw_file}" $((offset + SIZE_OFFSET)))"
  local size_dec
  size_dec=$((0x${size_hex}))
  if [[ "$size_dec" -eq 0 ]]; then
    die "Invalid image size: 0x${size_hex} (${size_dec})."
  fi

  info "Found cookies at ${offset}. Image size: 0x${size_hex} (${size_dec})."
  truncate -s "${size_dec}" "${ecrw_file}"
  info "File ${ecrw_file} truncated to ${size_dec} bytes."
}

extract_ecrw_files_from_ec() {
  local ecrw_file=$1
  local ecrw_ver_file=$2
  local ec_file=$3

  if ! futility dump_fmap -x "${ec_file}" "RW_FW:${ecrw_file}" >/dev/null 2>&1 ; then
    info "Falling back to EC_RW section for legacy EC."
    futility dump_fmap -x "${ec_file}" "EC_RW:${ecrw_file}" >/dev/null
    truncate_ec_rw "${ecrw_file}"
  fi
  info "EC RW extracted to ${ecrw_file}"

  futility dump_fmap -x "${ec_file}" "RW_FWID:${ecrw_ver_file}" >/dev/null
}

swap_ecrw() {
  local ap_file=$1
  local ec_file=$2
  local ec_config_file=$3
  local ap_for_ec_file=$4
  local raw_ecrw=$5
  local ec_ver=$6

  local temp_dir
  local ecrw_file
  local ecrw_hash_file
  local ecrw_ver_file
  local ecrw_config_file

  local region
  local info
  local ecrw_comp_type

  local ecrw_ver
  local apro_ver
  local aprw_ver

  temp_dir=$(mktemp -d)
  ecrw_file="${temp_dir}/${CBFS_ECRW_NAME}"
  ecrw_hash_file="${temp_dir}/${CBFS_ECRW_HASH_NAME}"
  ecrw_ver_file="${temp_dir}/${CBFS_ECRW_VERSION_NAME}"
  ecrw_config_file="${temp_dir}/${CBFS_ECRW_CONFIG_NAME}"

  if [[ -n "${ec_file}" ]]; then
    extract_ecrw_files_from_ec \
      "${ecrw_file}" \
      "${ecrw_ver_file}" \
      "${ec_file}"
    if [[ -n "${ec_config_file}" ]]; then
      ecrw_config_file="${ec_config_file}"
    fi
  elif [[ -n "${ap_for_ec_file}" ]]; then
    extract_ecrw_files_from_ap \
      "${ecrw_file}" \
      "${ecrw_ver_file}" \
      "${ecrw_config_file}" \
      "${ap_for_ec_file}"
  else
    ecrw_file="${raw_ecrw}"
    ecrw_ver_file="${ec_ver}"
  fi

  openssl dgst -sha256 -binary "${ecrw_file}" > "${ecrw_hash_file}"
  info "EC RW hash is recalculated and saved to ${ecrw_hash_file}"

  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_try_remove "${ap_file}" "${region}" "${CBFS_ECRW_VERSION_NAME}" \
      || true
    cbfstool_try_remove "${ap_file}" "${region}" "${CBFS_ECRW_CONFIG_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}"
    if [[ -e "${ecrw_ver_file}" ]]; then
      cbfstool "${ap_file}" add -r "${region}" -t raw \
        -c none -f "${ecrw_ver_file}" -n "${CBFS_ECRW_VERSION_NAME}"
    else
      warn "${CBFS_ECRW_VERSION_NAME} is missing from source file."
    fi
    # Add ecrw.config if provided.
    if [[ -e "${ecrw_config_file}" ]]; then
      cbfstool "${ap_file}" add -r "${region}" -t raw \
        -c "${ecrw_comp_type}" -f "${ecrw_config_file}" \
        -n "${CBFS_ECRW_CONFIG_NAME}"
    else
      warn "${CBFS_ECRW_CONFIG_NAME} is missing from source file."
    fi
  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}"

  if [[ -n "${ec_file}" ]]; then
    ecrw_ver=$(futility update --manifest -e "${ec_file}" \
      | jq -r '.default.ec.versions.rw')
  else
    # As some old `ap_for_ec_file` image may not have `ecrw_ver_file`, use
    # the `ap_for_ec_file` AP version as the EC RW version.
    ecrw_ver=$(futility update --manifest -i "${ap_for_ec_file}" \
      | jq -r '.default.host.versions.ro')
  fi
  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 are required."
  fi

  has_ec=$([[ -n "${FLAGS_ec}" ]] && echo 1 || echo 0)
  has_ap_for_ec=$([[ -n "${FLAGS_ap_for_ec}" ]] && echo 1 || echo 0)
  has_raw_ecrw=$([[ -n "${FLAGS_raw_ecrw}" ]] && echo 1 || echo 0)

  if [[ $((has_ec + has_ap_for_ec + has_raw_ecrw)) -ne 1 ]]; then
    flags_help
    die "Exactly one of -e/--ec, -a/--ap_for_ec, or -r/--raw_ecrw is required."
  fi

  if [[ -n "${FLAGS_ec}" ]] &&
     [[ -z "${FLAGS_ec_config}" ]] &&
     [[ -f "$(dirname ${FLAGS_ec})/ec.config" ]]; then
    FLAGS_ec_config="$(dirname ${FLAGS_ec})/ec.config"
    info "Using ec.config from ${FLAGS_ec_config}"
  fi
  if [[ -n "${FLAGS_ap_for_ec}" ]] &&
     [[ -n "${FLAGS_ec_config}" ]]; then
    die "-a/--ap_for_ec conflicts with --ec_config."
  fi
  if [[ -n "${FLAGS_ec}" ]] &&
     [[ -n "${FLAGS_ec_ver}" ]]; then
    die "-e/--ec conflicts with --ec_ver."
  fi
  if [[ -n "${FLAGS_ap_for_ec}" ]] &&
     [[ -n "${FLAGS_ec_ver}" ]]; then
    die "-a/--ap_for_ec conflicts with --ec_ver."
  fi

  swap_ecrw "${FLAGS_image}" "${FLAGS_ec}" "${FLAGS_ec_config}" "${FLAGS_ap_for_ec}" \
      "${FLAGS_raw_ecrw}" "${FLAGS_ec_ver}"
}

main "$@"
