#!/bin/bash

# Copyright (c) 2011 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 utility writes a firmware image to a tegra based board.

# Include common CrOS utilities.
. "/usr/lib/crosutils/common.sh"

get_default_board

# Command line options
DEFINE_string board "${DEFAULT_BOARD}" "The name of the board to flash."
DEFINE_string variant "" "Board variant."
DEFINE_string firmware "" "Path to firmware image to write."
DEFINE_string flasher "" "Path to flasher image to use."
DEFINE_string script "" "Path to flasher script to use."
DEFINE_string bct "" "Path to bct image to write."
DEFINE_string map "" "Path to System.map file to read for load address."
DEFINE_string config "" "Directory for temporary configuration files."
DEFINE_boolean sign $FLAGS_FALSE "Sign and append a BCT to the firmware image."
DEFINE_boolean dry_run $FLAGS_FALSE "Do everything except calling nvflash."

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

# Die on errors.
set -e

CROS_LOG_PREFIX="cros_write_firmware"

if [ -z "${FLAGS_board}" ]; then
  error "--board required."
  exit 1
fi

get_board_and_variant $FLAGS_board $FLAGS_variant

# Display the U-Boot version string.
function show_u_boot_info() {
  local firmware=$1
  local header=$(strings "${firmware}" | grep "U-Boot" | head -1 )

  info "Using ${header}"
}

# Given two values in decimal format return the first value rounded up to a
# multiple of the second value.  The second value must be a power of two.
function round() {
  awk "BEGIN { print and(($1 + $2 - 1), compl($2 - 1)) }"
}

# Given a value in decimal format return that value converted to hexidecimal
# and formatted as 0xXXXXXXXX.
function to_hex() {
  awk "BEGIN { printf \"0x%x\", $1 }"
}

# Create the flashing script.  This script will be loaded and run by the flasher
# build of U-Boot.  Most of the script is constant, we just need to append the
# sizes and address of the payload.
#
# TODO(robotboy): Compute the address from the text_base (0xe08000) +
# flasher image size (0x40000) + script block (0x1000)
function make_flashing_script() {
  local firmware=$1
  local script=$2
  local firmware_size=$(stat -c"%s" "${firmware}")

  echo "setenv address       0xe49000"
  echo "setenv firmware_size $(to_hex ${firmware_size})"
  echo "setenv length        $(to_hex $(round ${firmware_size} 4096))"

  cat "${script}"
}

# Parse the System.map file generated by the U-Boot build to get the entry
# point address.
function get_text_base() {
  local map=$1

  grep -m1 -E "^[0-9a-fA-F]{8} T _start$" ${map} | cut -d " " -f 1
}

###############################################################################

u_boot="/build/${BOARD_VARIANT}/u-boot"

# Generate default values for the firmware, flasher, BCT, map and flash config
# file paths from the board and variant information provided.
firmware=${FLAGS_firmware:-"${u_boot}/image.bin"}
flasher=${FLAGS_flasher:-"${u_boot}/u-boot-flasher.bin"}
script=${FLAGS_script:-"${u_boot}/flasher.script"}
bct=${FLAGS_bct:-"${u_boot}/bct/board.bct"}
map=${FLAGS_map:-"${u_boot}/System.map"}

# Extract the load and entry point for U-Boot.
text_base="0x$(get_text_base "${map}")"

# Verify that all of the files and tools we will need are present.
check_for_file "firmware" " " "${firmware}"
check_for_file "flasher" "  " "${flasher}"
check_for_file "script" "   " "${script}"
check_for_file "BCT" "      " "${bct}"
check_for_file "map" "      " "${map}"

check_for_tool "nvflash" "nvflash"

show_u_boot_info "${firmware}"

# If the user has not specified a config directory then create one on the fly
# and remove it when we are done.
if [ -z "${FLAGS_config}" ]; then
  config=$(mktemp -d)

  trap "rm -rf ${config}" EXIT
else
  mkdir -p "${FLAGS_config}"

  config=${FLAGS_config}
fi

# Optionally sign the firmware.  This is usefull if you want to write a bare
# U-Boot to flash and boot with that instead of a fully build ChromiumOS
# firmware image.
if [ $FLAGS_sign -eq $FLAGS_TRUE ]; then
  signed_firmware=${config}/signed.bin

  cros_sign_bootstub \
    --bct "${bct}" \
    --bootstub "${firmware}" \
    --output "${signed_firmware}" \
    --config "${config}" \
    --text_base "${text_base}"
else
  signed_firmware=${firmware}
fi

# Create the flashing script and turn it into a U-Boot uImage.
make_flashing_script "${signed_firmware}" "${script}" > ${config}/script

mkimage -A arm -T script -d ${config}/script ${config}/script.uimg > /dev/null

# Construct RAM image of flasher + script + firmware.  The firmware should
# already be signed by cbootimage.  The RAM image has a fixed 256K partition
# for the flasher, followed by the script and payload to be written.
dd if="${flasher}"            of="${config}/image"          &> /dev/null
dd if="${config}/script.uimg" of="${config}/image" seek=512 &> /dev/null
dd if="${signed_firmware}"    of="${config}/image" seek=520 &> /dev/null

if [ $FLAGS_dry_run -eq $FLAGS_FALSE ]; then
  # Write the RAM image image to the board.  The flasher will write the payload
  # into the boot device using the flashing script.
  sudo nvflash \
    --bct "${bct}" \
    --setbct \
    --bl "${config}/image" \
    --go \
    --setentry "${text_base}" "${text_base}" ||
    (echo && die "Failed to download U-Boot flasher + payload using nvflash.")

  # TODO(robotboy): Make a U-Boot tool that can compute the same checksum on
  # the host and display it here.  Unfortunately U-Boot uses a slightly
  # different crc32 than zlib, so it's not possible to use zlib's crc32 to do
  # this.  It has to be done by building a tool that uses U-Boots version.
  info "The serial terminal or LCD panel should show before and after CRCs of"
  info "the image written to the boot flash.  If these values match each"
  info "other your flash was probably successful."
fi