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

. /usr/share/misc/shflags

readonly DEFAULT_RETRIES=${DEFAULT_RETRIES:-4}
readonly STM32MON_CONNECT_RETRIES=${STM32MON_CONNECT_RETRIES:-6}
readonly STM32MON_SERIAL_BAUDRATE=${STM32MON_SERIAL_BAUDRATE:-115200}

DEFINE_boolean 'read' "${FLAGS_FALSE}" 'Read instead of write' 'r'
# Both flash read and write protection are removed by default, but you
# can optionally enable them (for testing purposes).
DEFINE_boolean 'remove_flash_read_protect' "${FLAGS_TRUE}" \
  'Remove flash read protection while performing command' 'U'
DEFINE_boolean 'remove_flash_write_protect' "${FLAGS_TRUE}" \
  'Remove flash write protection while performing command' 'u'
DEFINE_integer 'retries' "${DEFAULT_RETRIES}" 'Specify number of retries' 'R'
DEFINE_integer 'baudrate' "${STM32MON_SERIAL_BAUDRATE}" 'Specify UART baudrate' 'B'
DEFINE_boolean 'hello' "${FLAGS_FALSE}" 'Only ping the bootloader' 'H'
FLAGS_HELP="Usage: ${0} [flags] ec.bin"

# Process commandline flags
FLAGS "${@}" || exit 1
eval set -- "${FLAGS_ARGV}"

if [[ "$#" -eq 0 ]] && [[ "${FLAGS_hello}" -eq "${FLAGS_FALSE}" ]]; then
  echo "Missing filename"
  flags_help
  exit 1
fi

# print out canonical path to differentiate between /usr/local/bin and
# /usr/bin installs
echo "$(readlink -f "$0")"

readonly CROS_EC_SPI_MODALIAS_STR="of:NcrfpTCgoogle,cros-ec-spi"

readonly CROS_EC_UART_MODALIAS_STR="of:NcrfpTCgoogle,cros-ec-uart"

check_hardware_write_protect_disabled() {
  local hardware_write_protect_state="$(crossystem wpsw_cur)"
  if [[ "${hardware_write_protect_state}" != "0" ]]; then
    echo "Please make sure hardware write protect is disabled."
    echo "See https://www.chromium.org/chromium-os/firmware-porting-guide/firmware-ec-write-protection"
    exit 1
  fi
}

# Get the spiid for the fingerprint sensor based on the modalias
# string: https://crbug.com/955117
get_spiid() {
  # TODO(b/179533783): Fix modalias on strongbad and remove this bypass.
  if [[ -n "${DEVICEID}" ]]; then
    echo "${DEVICEID}"
    exit 0
  fi

  for dev in /sys/bus/spi/devices/*; do
    if [[ "$(cat "${dev}/modalias")" == "${CROS_EC_SPI_MODALIAS_STR}" ]]; then
      echo "$(basename "${dev}")"
      exit 0
    fi
  done

  exit 1
}

# Get the uartid for the fingerprint sensor based on the modalias
get_uartid() {
  for dev in /sys/bus/serial/devices/*; do
    if [ -f "${dev}/modalias" ]; then
      if [[ "$(cat "${dev}/modalias")" == "${CROS_EC_UART_MODALIAS_STR}" ]]; then
        echo "$(basename "${dev}")"
        exit 0
      fi
    fi
  done

  exit 1
}

# Usage: gpio <unexport|export|in|out|0|1|get> <signal> [signal...]
gpio() {
  local cmd="$1"
  shift

  for signal in "$@"; do
    case "${cmd}" in
      unexport|export)
        echo "${signal}" > "/sys/class/gpio/${cmd}"
        ;;
      in|out)
        local direction="${cmd}"
        echo "${direction}" > "/sys/class/gpio/gpio${signal}/direction"
        ;;
      0|1)
        local value="${cmd}"
        echo "${value}" > "/sys/class/gpio/gpio${signal}/value"
        ;;
      get)
        local value="${cmd}"
        cat "/sys/class/gpio/gpio${signal}/value"
        ;;
      *)
        echo "Invalid gpio command: ${cmd}" >&2
        exit 1
        ;;
    esac
  done
}

# Taken verbatim from
# https://chromium.googlesource.com/chromiumos/docs/+/master/lsb-release.md#shell
# This should not be used by anything except get_platform_name.
# See https://crbug.com/98462.
lsbval() {
  local key="$1"
  local lsbfile="${2:-/etc/lsb-release}"

  if ! echo "${key}" | grep -Eq '^[a-zA-Z0-9_]+$'; then
    return 1
  fi

  sed -E -n -e \
    "/^[[:space:]]*${key}[[:space:]]*=/{
      s:^[^=]+=[[:space:]]*::
      s:[[:space:]]+$::
      p
    }" "${lsbfile}"
}

# Get the underlying board (reference design) that we're running on (not the
# FPMCU or sensor).
get_platform_name() {
  local platform_name

  # We used to use "cros_config /identity platform-name", but that is specific
  # to mosys and does not actually provide the board name in all cases.
  # cros_config intentionally does not provide a way to get the board
  # name: b/156650654.

  # If there was a way to get the board name from cros_config, it's possible
  # that it could fail in the following cases:
  #
  # 1) We're running on a non-unibuild device (the only one with FP is nocturne)
  # 2) We're running on a proto device during bringup and the cros_config
  #    settings haven't yet been setup.
  #
  # In all cases we can fall back to /etc/lsb-release. It's not recommended
  # to do this, but we don't have any other options in this case.
  echo "Getting platform name from /etc/lsb-release." 1>&2
  platform_name="$(lsbval "CHROMEOS_RELEASE_BOARD")"
  if [[ -z "${platform_name}" ]]; then
    exit 1
  fi

  echo "${platform_name}"
}

check_gpio_chip_exists() {
  local gpiochip="$1"
  if [[ ! -e "/sys/class/gpio/${gpiochip}" ]]; then
    echo "Cannot find GPIO chip: ${gpiochip}"
    exit 1
  fi
}

flash_fp_mcu_stm32() {
  local transport="${1}"
  local device="${2}"
  local gpio_nrst="${3}"
  local gpio_boot0="${4}"
  local gpio_pwren="${5}"
  local file="${6}"
  local deviceid

  local stm32mon_flags="-p --retries ${STM32MON_CONNECT_RETRIES}"

  if [[ "${transport}" == "UART" ]]; then
    stm32mon_flags+=" --baudrate ${FLAGS_baudrate} --device ${device}"
  else
    stm32mon_flags+=" -s ${device}"
  fi

  if [[ "${FLAGS_hello}" -eq "${FLAGS_FALSE}" ]]; then
    if [[ "${FLAGS_remove_flash_write_protect}" -eq "${FLAGS_TRUE}" ]]; then
      stm32mon_flags+=" -u"
    fi

    if [[ "${FLAGS_remove_flash_read_protect}" -eq "${FLAGS_TRUE}" ]]; then
      stm32mon_flags+=" -U"
    fi

    if [[ "${FLAGS_read}" -eq "${FLAGS_TRUE}" ]]; then
      # Read from FPMCU to file
      if [[ -e "${file}" ]]; then
        echo "Output file already exists: ${file}"
        exit 1
      fi
      stm32mon_flags+=" -r ${file}"
    else
      # Write to FPMCU from file
      if [[ ! -f "${file}" ]]; then
        echo "Invalid image file: ${file}"
        exit 1
      fi
      stm32mon_flags+=" -e -w ${file}"
    fi
  fi


  check_hardware_write_protect_disabled

  if [[ "${transport}" == "UART" ]]; then
    deviceid="$(get_uartid)"
    if [[ $? -ne 0 ]]; then
      echo "Unable to find FP sensor UART device: ${CROS_EC_UART_MODALIAS_STR}"
      exit 1
    fi
  else
    deviceid="$(get_spiid)"
    if [[ $? -ne 0 ]]; then
      echo "Unable to find FP sensor SPI device: ${CROS_EC_SPI_MODALIAS_STR}"
      exit 1
    fi
  fi

  echo "Flashing ${transport} device ID: ${deviceid}"

  # Ensure the ACPI is not cutting power when unloading cros-ec-spi
  if [[ "${gpio_pwren}" -gt 0 ]]; then
    gpio export "${gpio_pwren}"
    gpio out    "${gpio_pwren}"
    gpio 1      "${gpio_pwren}"
  fi

  # Remove cros_fp if present
  if [[ "${transport}" == "UART" ]]; then
    echo "${deviceid}" > /sys/bus/serial/drivers/cros-ec-uart/unbind
  else
    echo "${deviceid}" > /sys/bus/spi/drivers/cros-ec-spi/unbind
  fi

  # Configure the MCU Boot0 and NRST GPIOs
  gpio export "${gpio_boot0}" "${gpio_nrst}"
  gpio out    "${gpio_boot0}" "${gpio_nrst}"

  # Reset sequence to enter bootloader mode
  gpio 1 "${gpio_boot0}"
  gpio 0 "${gpio_nrst}"
  sleep 0.001

  if [[ "${transport}" == "UART" ]]; then
    # load AMDI0020:01 ttyS1
    echo AMDI0020:01 > /sys/bus/platform/drivers/dw-apb-uart/unbind;
    echo AMDI0020:01 > /sys/bus/platform/drivers/dw-apb-uart/bind;
  else
    # load spidev (fail on cros-ec-spi first to change modalias)
    echo "${deviceid}" > /sys/bus/spi/drivers/cros-ec-spi/bind 2>/dev/null
    echo "${deviceid}" > /sys/bus/spi/drivers/spidev/bind
    # The following sleep is a workaround to mitigate the effects of a
    # poorly behaved chip select line. See b/145023809.
  fi
  sleep 0.5

  # We do not expect the drivers to change the pin state when binding.
  # If you receive this warning, the driver needs to be fixed on this board
  # and this flash attempt will probably fail.
  if [[ $(gpio get "${gpio_boot0}") -ne 1 ]]; then
    echo "WARNING: One of the drivers changed BOOT0 pin state on bind attempt."
  fi
  if [[ $(gpio get "${gpio_nrst}") -ne 0 ]]; then
    echo "WARNING: One of the drivers changed NRST pin state on bind attempt."
  fi

  # TODO(b/179530529): Remove this hack when the drivers on strongbad stop
  # setting the gpio states.
  if [[ "${PLATFORM_NAME}" == "strongbad" ]]; then
    echo "# Fixing GPIOs"
    gpio 1 "${gpio_boot0}"
    gpio 0 "${gpio_nrst}"
  fi

  local attempt=0
  local cmd_exit_status=1
  local cmd="stm32mon ${stm32mon_flags}"

  for attempt in $(seq ${FLAGS_retries}); do
    # Reset sequence to enter bootloader mode
    gpio 0 "${gpio_nrst}"
    sleep 0.01

    # Release reset as the SPI bus is now ready
    gpio 1 "${gpio_nrst}"

    # As per section '68: Bootloader timings' from application note below:
    # https://www.st.com/resource/en/application_note/cd00167594-stm32-microcontroller-system-memory-boot-mode-stmicroelectronics.pdf
    # bootloader startup time is 16.63 ms for STM32F74xxx/75xxx and 53.975 ms
    # for STM32H74xxx/75xxx. SPI needs 1 us delay for one SPI byte sending.
    # Keeping some margin, add delay of 100 ms to consider minimum bootloader
    # startup time after the reset for stm32 devices.
    sleep 0.1

    # Print out the actual underlying command we're running and run it
    echo "# ${cmd}"
    ${cmd}
    cmd_exit_status=$?

    if [[ "${cmd_exit_status}" -eq 0 ]]; then
      break
    fi
    echo "# Attempt ${attempt} failed."
    echo
    sleep 1
  done

  # unload device
  if [[ "${transport}" != "UART" ]]; then
    echo "${deviceid}" > /sys/bus/spi/drivers/spidev/unbind
  fi

  # Go back to normal mode
  gpio out "${gpio_nrst}"
  gpio 0   "${gpio_boot0}" "${gpio_nrst}"
  gpio 1   "${gpio_nrst}"

  # Give up GPIO control
  # TODO(b/179839337): Remove this for Strongbad DVT, where we should have
  # correct external pulls on these lines.
  if [[ "${PLATFORM_NAME}" != "strongbad" ]]; then
    gpio in       "${gpio_boot0}" "${gpio_nrst}"
  fi
  gpio unexport "${gpio_boot0}" "${gpio_nrst}"

  # wait for FP MCU to come back up (including RWSIG delay)
  sleep 2

  # Put back cros_fp driver if transport is SPI
  if [[ "${transport}" != "UART" ]]; then
    echo "${deviceid}" > /sys/bus/spi/drivers/cros-ec-spi/bind

    # Kernel driver is back, we are no longer controlling power
    if [[ "${gpio_pwren}" -gt 0 ]]; then
      gpio unexport "${gpio_pwren}"
    fi
  fi

  if [[ "${cmd_exit_status}" -ne 0 ]]; then
    exit 1
  fi

  # Inform user to reboot if transport is UART.
  # Display fw version is transport is SPI
  if [[ "${transport}" == "UART" ]]; then
    echo "Please reboot this device."
  else
    # Test it
    ectool --name=cros_fp version
  fi
}

config_hatch() {
  check_gpio_chip_exists "gpiochip200"

  readonly TRANSPORT="SPI"
  readonly DEVICE="/dev/spidev1.1"
  # See
  # third_party/coreboot/src/soc/intel/cannonlake/include/soc/gpio_soc_defs.h
  # for pin name to number mapping.
  # Examine `cat /sys/kernel/debug/pinctrl/INT34BB:00/gpio-ranges` on a hatch
  # device to determine gpio number from pin number.
  # FPMCU RST_ODL is on GPP_A12 = 200 + 12 = 212
  readonly GPIO_NRST=212
  # FPMCU BOOT0 is on GPP_A22 = 200 + 22 = 222
  readonly GPIO_BOOT0=222
  # FP_PWR_EN is on GPP_C11 = 456 + (192 - 181) = 456 + 11 = 467
  readonly GPIO_PWREN=467
}


config_nami() {
  check_gpio_chip_exists "gpiochip360"

  readonly TRANSPORT="SPI"
  readonly DEVICE="/dev/spidev32765.0"

  # FPMCU RST_ODL is on GPP_C9 = 360 + 57 = 417
  readonly GPIO_NRST=417
  # FPMCU BOOT0 is on GPP_D5 = 360 + 77 = 437
  readonly GPIO_BOOT0=437
  # FP_PWR_EN is on GPP_B11 = 360 + 35 = 395
  readonly GPIO_PWREN=395
}

config_nocturne() {
  check_gpio_chip_exists "gpiochip360"

  readonly TRANSPORT="SPI"
  readonly DEVICE="/dev/spidev32765.0"

  # FPMCU RST_ODL is on GPP_C10 = 360 + 58 = 418
  readonly GPIO_NRST=418
  # FPMCU BOOT0 is on GPP_C8 = 360 + 56 = 416
  readonly GPIO_BOOT0=416
  # FP_PWR_EN is on GPP_A11 = 360 + 11 = 371
  readonly GPIO_PWREN=371
}

config_strongbad() {
  check_gpio_chip_exists "gpiochip392"

  readonly TRANSPORT="SPI"
  readonly DEVICE="/dev/spidev10.0"
  # TODO(b/179533783): Fix modalias on strongbad and remove this bypass.
  readonly DEVICEID="spi10.0"

  # FPMCU RST_ODL is $(gpiofind FP_RST_L) is gpiochip0 22
  readonly GPIO_NRST=$((392 + $(gpiofind FP_RST_L|cut -f2 -d" ")))
  # FPMCU BOOT0 is $(gpiofind FPMCU_BOOT0) is gpiochip0 10
  readonly GPIO_BOOT0=$((392 + $(gpiofind FPMCU_BOOT0|cut -f2 -d" ")))
  # TODO(b/179839337): Hardware currently doesn't support PWREN, but the
  # next revision will. Add a comment here about the power enable gpio.
  readonly GPIO_PWREN=-1
}

config_volteer() {
  check_gpio_chip_exists "gpiochip152"

  readonly TRANSPORT="SPI"
  readonly DEVICE="/dev/spidev1.0"

  # See kernel/v5.4/drivers/pinctrl/intel/pinctrl-tigerlake.c
  # for pin name and pin number.
  # Examine `cat /sys/kernel/debug/pinctrl/INT34C5:00/gpio-ranges` on a
  # volteer device to determine gpio number from pin number.
  # For example: GPP_C23 is UART2_CTS which can be queried from EDS
  # the pin number is 194. From the gpio-ranges, the gpio value is
  # 408 + (194-171) = 431

  # FPMCU RST_ODL is on GPP_C23 = 408 + (194 - 171) = 431
  readonly GPIO_NRST=431
  # FPMCU BOOT0 is on GPP_C22 = 408 + (193 - 171) = 430
  readonly GPIO_BOOT0=430
  # FP_PWR_EN is on GPP_A21 = 216 + (63 - 42) = 237
  readonly GPIO_PWREN=237
}

config_zork() {
  check_gpio_chip_exists "gpiochip256"

  readonly TRANSPORT="UART"
  readonly DEVICE="/dev/ttyS1"

  # FPMCU RST_ODL is on AGPIO 11 = 256 + 11 = 267
  readonly GPIO_NRST=267
  # FPMCU BOOT0 is on AGPIO 69 = 256 + 69 = 325
  readonly GPIO_BOOT0=325
  # FPMCU PWR_EN is on AGPIO 32 = 256 + 32 = 288, but should not be
  # necessary for flashing. Set invalid value.
  readonly GPIO_PWREN=-1
}

config_guybrush() {
  check_gpio_chip_exists "gpiochip256"

  readonly TRANSPORT="UART"

  readonly DEVICE="/dev/ttyS1"

  # FPMCU RST_ODL is on AGPIO 11 = 256 + 11 = 267
  readonly GPIO_NRST=267
  # FPMCU BOOT0 is on AGPIO 144 = 256 + 144 = 400
  readonly GPIO_BOOT0=400
  # FPMCU PWR_EN is on AGPIO 32 = 256 + 32 = 288, but should not be
  # necessary for flashing. Set invalid value.
  readonly GPIO_PWREN=-1
}

# The "platform name" corresponds to the underlying board (reference design)
# that we're running on (not the FPMCU or sensor). At the moment all of the
# reference designs use the same GPIOs. If for some reason a design differs in
# the future, we will want to add a nested check in the config_<platform_name>
# function. Doing it in this manner allows us to reduce the number of
# configurations that we have to maintain (and reduces the amount of testing
# if we're only updating a specific config_<platform_name>).
readonly PLATFORM_NAME="$(get_platform_name)"

echo "Using config for ${PLATFORM_NAME}"

# Check that the config function exists
if [[ "$(type -t "config_${PLATFORM_NAME}")" != "function" ]]; then
  echo "No config for platform ${PLATFORM_NAME}"
  exit 1
fi

config_${PLATFORM_NAME}

if [[ $? -ne 0 ]]; then
  echo "Configuration failed for platform ${PLATFORM_NAME}"
  exit 1
fi

flash_fp_mcu_stm32 \
  "${TRANSPORT}"   \
  "${DEVICE}"      \
  "${GPIO_NRST}"   \
  "${GPIO_BOOT0}"  \
  "${GPIO_PWREN}"  \
  "${1}"