#!/bin/bash

#
# Copyright 2016 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.
#
# This script is a utility which allows to create differently signed CR50
# images from different sources.
#
set -e
set -u

progname=$(basename $0)

OD="/usr/bin/od"

tmpf="/tmp/bs_manifest.$$"
trap "{ if [[ -f 1.flat ]]; then rm -rf [01].flat ${tmpf} ; fi; }" EXIT

usage() {
  local rv="${1}"
  cat <<EOF

This script allows to sign CR50 RW images. By default it uses ec.RW.elf and
ec.RW_B.elf in build/cr50/RW as inputs and util/signer/ec_RW-manifest-dev.json
as the manifest, and places the newly signed images into build/cr50/ec.bin.

The only outside dependency of this script is the signing utility itself,
which is expected to be available as /usr/bin/cr50-codesigner.

The utility can be installed by running 'sudo emerge cr50-utils',

The following command line options are accepted:

  b1 - generate signature for the b1 version of the H1 chip
  elves <elf1> <elf2> - sign the supplied elf files instead of the default
        ones. Handy if the builder generated files need to be signed
  help - print this message
  hex - generate hex output instead of binary (place in 0.signed.hex and
        1.signed.hex in the local directory)
  prod - sign with prod key (no debug image will be signed)

This script also allows to sign dev images for running on prod RO. To do that
invoke this script as follows:

  H1_DEVIDS='<dev id0> <dev id1>' ${progname} [other options, if any]

where <dev id0> <dev id1> are values reported by sysinfo command in the
DEV_ID: line when run on the CR50 for which the image is built.

The same values can be obtained in the lsusb command output:

  lsusb -vd 18d1:5014 | grep -i serial

note that the lsusb reported values are in hex and need to be prefixed with
0x.

Finally, this script also allows to specify the board ID fields of the RW
headers. The fields come from the evironment variable CR50_BOARD_ID, which is
required to include three colon separated fields. The first field is a four
letter board RLZ code, the second field is board id mask in hex, no 0x prefix,
and the third field - board ID flags, again, hex, no 0x prefix.

CR50_BOARD_ID='XXYY:ffffff00:ff00' ${progname} [other options, if any]

both H1_DEVIDS and CR50_BOARD_ID can be defined independently.

EOF
  exit "${rv}"
}

# This function modifies the manifest to include device ID and board ID nodes,
# if H1_DEVIDS and CR50_BOARD_ID are defined in the environment, respectively,
tweak_manifest () {
  local sub

  # If defined, plug in dev ID nodes before the 'fuses' node.
  if [[ -z "${do_prod}" && -n "${H1_DEVIDS}" ]]; then
    echo "creating a customized DEV image for DEV IDS ${H1_DEVIDS}"
    sub=$(printf "\\\n   \"DEV_ID0\": %s,\\\n   \"DEV_ID1\": %s," ${H1_DEVIDS})
    sed -i "s/\"fuses\": {/\"fuses\": {${sub}/" "${tmpf}"
  fi

  if [[ -z "${CR50_BOARD_ID}" ]]; then
    return
  fi

  # CR50_BOARD_ID is set, let's parse it and plug in the board ID related
  # nodes into manifest before the 'fuses' node.
  local bid_params
  local rlz

  bid_params=( $(echo $CR50_BOARD_ID | sed 's/:/ /g') )
  # A very basic sanity check: it needs to consist of three colon separated
  # fields.
  if [[ ${#bid_params[@]} != 3 ]]; then
    echo "Wrong board ID string \"$CR50_BOARD_ID\"}" >&2
    exit 1
  fi

  if [[ "${bid_params[0]}" == "0" ]] ; then
    rlz="0"
  elif [[ ${#bid_params[0]} == 4 ]] ; then
    # Convert 4 char board RLZ code from ASCII to hex
    rlz="0x$(echo -n ${bid_params[0]} | hexdump -ve '/1 "%02x"')"
  else
    echo "Invalid RLZ ${bid_params[0]}"
    exit 1
  fi

  # Prepare text of all three board ID related nodes
  sub="$(printf "\\\n\"board_id\": %s,\\\n" "${rlz}")"
  sub+="$(printf "\"board_id_mask\": %s,\\\n" "0x${bid_params[1]}")"
  sub+="$(printf "\"board_id_flags\": %s,\\\n" "0x${bid_params[2]}")"
  sed -i "s/\"fuses\": {/${sub}\"fuses\": {/" "${tmpf}"
}

# This function accepts two arguments, names of two binary files.
#
# It searches the first passed in file for the first 8 bytes of the second
# passed in file. The od utility is used to generate full hex dump of the
# first file (16 bytes per line) and the first 8 bytes of the second file.
#
# grep is used to check if the pattern is present in the full dump. If the
# pattern is not found, the first file is dumped again, this time with an 8
# byte offset into the file. This makes sure that if the match is present, but
# is spanning two lines of the original hex dump, it is in a single dump line
# the second time around.
find_blob_in_blob() {
  local main_blob="${1}"
  local pattern_blob="${2}"
  local pattern
  local od_options="-An -tx1"

  # Get the first 8 bytes of the pattern blob.
  pattern="$(${OD} ${od_options} -N8 "${pattern_blob}")"

  if "${OD}" ${od_options} "${main_blob}" | grep "${pattern}" > /dev/null; then
    return 0
  fi

  # Just in case pattern was wrapped in the previous od output, let's do it
  # again with an 8 bytes offset
  if "${OD}" ${od_options} -j8 "${main_blob}" |
      grep "${pattern}" > /dev/null; then
    return 0
  fi

  return 1
}

# This function accepts two arguments, names of the two elf files.
#
# The files are searched for test RMA public key patterns - x25519 or p256,
# both files are supposed to have pattern of one of these keys and not the
# other. If this holds true the function prints the public key base name. If
# not both files include the same key, or include more than one key, the
# function reports failure and exits the script.
determine_rma_key_base() {
  local base_name="${EC_ROOT}/board/cr50/rma_key_blob"
  local curve
  local curves=( "x25519" "p256" )
  local elf
  local elves=( "$1" "$2" )
  local key_file
  local mask=1
  local result=0

  for curve in ${curves[@]}; do
    key_file="${base_name}.${curve}.test"
    for elf in ${elves[@]}; do
      if find_blob_in_blob "${elf}" "${key_file}"; then
        result=$(( result | mask ))
      fi
      mask=$(( mask << 1 ))
    done
  done

  case "${result}" in
    (3)  curve="x25519";;
    (12) curve="p256";;
    (*) echo "could not determine key type in the elves" >&2
        exit 1
        ;;
  esac

  echo "${base_name}.${curve}"
}

SIGNER="cr50-codesigner"
if ! which "${SIGNER}" 2>/dev/null > /dev/null; then
  echo "${SIGNER} is not available, try running 'sudo emerge cr50-utils'" >&2
  exit 1
fi

# This is where the new signed image will be pasted into.
: ${RESULT_FILE=build/cr50/ec.bin}
TMP_RESULT_FILE="${RESULT_FILE}.tmp"

if [[ -z "${CROS_WORKON_SRCROOT}" ]]; then
  echo "${progname}: This script must run inside Chrome OS chroot" >&2
  exit 1
fi

: ${CR50_BOARD_ID=}
: ${H1_DEVIDS=}
EC_ROOT="${CROS_WORKON_SRCROOT}/src/platform/ec"
EC_BIN_ROOT="${EC_ROOT}/util/signer"

do_hex=
do_b1=
do_prod=

# Prepare the default manifest.
cp "${EC_BIN_ROOT}/ec_RW-manifest-dev.json" "${tmpf}"

elves=( build/cr50/RW/ec.RW.elf build/cr50/RW/ec.RW_B.elf )
cd "${EC_ROOT}"
while (( $# )); do
  param="${1}"
  case "${param}" in
    (hex) do_hex='true';;
    (b1)
      do_b1='true'
      sed -i 's/\(.*FW_DEFINED_DATA_BLK0.*\): 2/\1: 0/' "${tmpf}"
      ;;
    (elves)
      if [[ (( $# < 3 )) ]]; then
        echo "two elf file names are required" >&2
        exit 1
      fi
      elves=( $2 $3 )
      shift
      shift
      ;;
    (prod)
      do_prod='true'
      ;;
    (help)
      usage 0
      ;;
    (*)
      usage 1
      ;;
  esac
  shift
done

if [[ -z "${do_hex}" && ! -f "${RESULT_FILE}" ]]; then
  echo "${RESULT_FILE} not found. Run 'make BOARD=cr50' first" >&2
  exit 1
fi

if [[ -n "${do_prod}" && -n "${do_b1}" ]]; then
  echo "can not build prod images for B1, sorry..."
  exit 1
fi

# If signing a chip factory image (version 0.0.22) do not try figuring out the
# RMA keys.
ignore_rma_keys="$(awk '
  BEGIN {count = 0};
   /"major": 0,/ {count += 1};
   /"minor": 22,/ {count += 1};
   END {{if (count == 2) {print "yes"};}}' \
        "${EC_BIN_ROOT}/ec_RW-manifest-prod.json")"

if [ "${ignore_rma_keys}" != "yes" ]; then
  rma_key_base="$(determine_rma_key_base ${elves[@]})"
else
  echo "Ignofing RMA keys for factory branch"
fi

signer_command_params=()
signer_command_params+=(--b -x ${EC_BIN_ROOT}/fuses.xml)
if [[ -z "${do_prod}" ]]; then
  signer_command_params+=(-k ${EC_BIN_ROOT}/cr50_rom0-dev-blsign.pem.pub)
else
  cp "${EC_BIN_ROOT}/ec_RW-manifest-prod.json" "${tmpf}"
  signer_command_params+=(-k ${EC_BIN_ROOT}/cr50_RW-prod.pem.pub)
  # Swap test public RMA server key with the prod version.
  if [ "${ignore_rma_keys}" != "yes" ]; then
    signer_command_params+=(-S "${rma_key_base}.test","${rma_key_base}.prod")
  fi
fi
signer_command_params+=(-j ${tmpf})

if [[ -n "${do_hex}" ]]; then
  dst_suffix='signed.hex'
else
  signer_command_params+=(--format=bin)
  dst_suffix='flat'
fi

tweak_manifest

count=0
for elf in ${elves[@]}; do
  if [[ -n "${do_prod}" ]]; then
    if strings "${elf}" | egrep -q "(DBG|SQA)/cr50"; then
      echo "Will not sign debug or SQA image with prod keys" >&2
      exit 1
    fi
  fi
  signed_file="${count}.${dst_suffix}"

  # Make sure output file is not owned by root
  touch "${signed_file}"
  command="${SIGNER} ${signer_command_params[@]} -i ${elf} -o ${signed_file}"
  if ! ${command}; then
    echo "${progname}: \"${command}\" failed" >&2
    exit 1
  fi

  if [ "${ignore_rma_keys}" != "yes" ]; then
    if find_blob_in_blob  "${signed_file}" "${rma_key_base}.test"; then
      echo "${progname}: test RMA key in the signed image!" >&2
      rm *."${dst_suffix}"
      exit 1
    fi

    if ! find_blob_in_blob "${signed_file}" "${rma_key_base}.prod"; then
      echo "${progname}: prod RMA key not in the signed image!" >&2
      rm *."${dst_suffix}"
      exit 1
    fi
  fi
  : $(( count++ ))
done

if [[ -z "${do_hex}" ]]; then
  # Full binary image is required, paste the newly signed blobs into the
  # output image, preserving it in case dd fails for whatever reason.
  cp "${RESULT_FILE}" "${TMP_RESULT_FILE}"
  dd if="0.flat" of="${TMP_RESULT_FILE}" seek=16384 bs=1 conv=notrunc
  dd if="1.flat" of="${TMP_RESULT_FILE}" seek=278528 bs=1 conv=notrunc
  rm [01].flat
  mv "${TMP_RESULT_FILE}" "${RESULT_FILE}"
fi

echo "SUCCESS!!!"