Rewrite flash script to behave like Enterprise

- Rewrite this like a real program, lots of things cribbed from the
Enterprise script.
- Add serial number parameter
- Allow specifying path of files, otherwise use PRODUCT_OUT, and finally
fall back to the script's directory.
- Leave in a flag to allow just booting LK (-l)

Change-Id: I26c997b345da00fc5efedb07c450c5d4fecb6c3b
diff --git a/flash.sh b/flash.sh
index 5bf78a3..666950c 100755
--- a/flash.sh
+++ b/flash.sh
@@ -14,140 +14,190 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-set -e
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+FASTBOOT_CMD="$(which fastboot)"
 
-SCRIPT_DIR=$(realpath $(dirname $0))
-PYTHON3=$(which python3)
-FLASHTOOLS_DIR=${SCRIPT_DIR}/flashtools
-FBTOOL="${PYTHON3} ${FLASHTOOLS_DIR}/fbtool.py"
-EXCELSIOR_DEBUG_RESET_TOOL="${PYTHON3} ${FLASHTOOLS_DIR}/debug_board_reset.py"
-DOWNLOAD_AGENT_CONFIG=${FLASHTOOLS_DIR}/dl_addr.ini
-DOWNLOAD_AGENT_CONFIG_PRODUCT_OUT=${PRODUCT_OUT}/dl_addr.ini
-FASTBOOT_TOOL=$(which fastboot)
-USERSPACE_ARCH=arm64
-
-# Check for fastboot tool
-if [[ -z "${FASTBOOT_TOOL}" ]]; then
-    echo "Fastboot is not found on this machine. Install fastboot before running this script"; exit 1
-fi
-
-# Check if PRODUCT_OUT is set
-if [[ -z "${PRODUCT_OUT}" ]]; then
-    echo "PRODUCT_OUT not set, attempting to use script path ${SCRIPT_DIR}.";
-    PRODUCT_OUT=${SCRIPT_DIR}
-    DOWNLOAD_AGENT_CONFIG_PRODUCT_OUT=${PRODUCT_OUT}/dl_addr.ini
-    # Check that download agent config is available.
-    if [[ ! -f "${DOWNLOAD_AGENT_CONFIG_PRODUCT_OUT}" ]]; then
-        echo "Download agent config not found in ${PRODUCT_OUT}."; exit 1
-    fi
-else
-    # Create symlink to download agent config.
-    ln -sf ${ROOTDIR}/board/flashtools/dl_addr.ini ${DOWNLOAD_AGENT_CONFIG_PRODUCT_OUT}
-fi
-
-
-show_help () {
-	echo "\
-Script for flashing excelsior
-Example: $0
-	Options:
-		-b	option to flash only boot partition."
+function die {
+    echo "$@" >/dev/stderr
+    exit 1
 }
 
-
-reset_excelsior () {
-  ${EXCELSIOR_DEBUG_RESET_TOOL} &
+function try {
+    $@ || die "Failed to execute '$@'"
 }
 
-reset_excelsior_to_rom() {
-  echo "Reset Excelsior to ROM"
-  LANG=en_US.UTF-8 LC_ALL=en_US.UTF-8 ${EXCELSIOR_DEBUG_RESET_TOOL} --rom &
-}
-
-load_lk_to_ram() {
-  pushd ${PRODUCT_OUT}
-  echo "Load LK to rom and boot LK"
-  ${FBTOOL} -f ${DOWNLOAD_AGENT_CONFIG_PRODUCT_OUT}
+function load_fastboot {
+  local files_dir="$1"
+  python3 "${files_dir}/flashtools/debug_board_reset.py" --rom &
+  pushd "${files_dir}"
+  python3 "${files_dir}/flashtools/fbtool.py" -f "${files_dir}/dl_addr.ini"
   popd
 }
 
-load_fastboot() {
-	reset_excelsior_to_rom
-	sleep 1
-	load_lk_to_ram
-	sleep 1
+function flash_partitions {
+  local files_dir="$1"
+  load_fastboot "${files_dir}"
+  try ${FASTBOOT_CMD} erase mmc0
+  try ${FASTBOOT_CMD} erase mmc0boot0
+  try ${FASTBOOT_CMD} erase mmc0boot1
+  try ${FASTBOOT_CMD} flash mmc0 "${files_dir}/partition-table.img"
+  try ${FASTBOOT_CMD} flash mmc0boot0 "${files_dir}/bl2.img"
+  try ${FASTBOOT_CMD} flash mmc0boot1 "${files_dir}/u-boot-env.bin"
+  try ${FASTBOOT_CMD} flash bootloaders "${files_dir}/fip.bin"
+  try ${FASTBOOT_CMD} flash boot "${files_dir}/boot_arm64.img"
+  try ${FASTBOOT_CMD} flash rootfs "${files_dir}/rootfs_arm64.img"
 }
 
-erase_mmc() {
-  echo "Erase EMMC"
-  ${FASTBOOT_TOOL} erase mmc0
-  ${FASTBOOT_TOOL} erase mmc0boot0
-  ${FASTBOOT_TOOL} erase mmc0boot1
+function ensure_flash_files_present {
+  local found_all_files=true
+  local files_dir="$1"; shift
+  local files=(
+    "bl2.img"
+    "boot_arm64.img"
+    "dl_addr.ini"
+    "fip.bin"
+    "flashtools/debug_board_reset.py"
+    "flashtools/dl_addr.ini"
+    "flashtools/fbtool.py"
+    "lk.bin"
+    "partition-table.img"
+    "rootfs_arm64.img"
+    "u-boot-env.bin"
+  )
+
+  if [[ "${PRODUCT_OUT}" == "${files_dir}" ]]; then
+    ln -sf "${ROOTDIR}/board/flashtools/dl_addr.ini" "${files_dir}/dl_addr.ini"
+    rm -f "${files_dir}/flashtools"
+    ln -sf "${ROOTDIR}/board/flashtools" "${files_dir}/flashtools"
+  fi
+
+  for file in "${files[@]}"; do
+    if [[ ! -f "${files_dir}/${file}" ]]; then
+      echo "${file} is missing!"
+      found_all_files=""
+    fi
+  done
+
+  if [[ -z "${found_all_files}" ]]; then
+    die "Required files are missing. Can not continue."
+  fi
 }
 
+function detect_device_or_die {
+  local retries="$1"; shift
+  local serial="$1"; shift
+  local count=0
+  local output
+  local found
 
-flash_fastboot() {
-    pushd ${PRODUCT_OUT}
-    erase_mmc
-    ${FASTBOOT_TOOL} flash mmc0 partition-table.img
-    ${FASTBOOT_TOOL} flash mmc0boot0 bl2.img
-    ${FASTBOOT_TOOL} flash mmc0boot1 u-boot-env.bin
-    ${FASTBOOT_TOOL} flash bootloaders fip.bin
-    popd
+  echo -n "Waiting up to ${retries} seconds for device ${serial}"
+  for ((count=0; count<"${retries}"; count++)); do
+    output=$(${FASTBOOT_CMD} devices)
+    if [[ "$?" != "0" ]]; then
+      die "Unable to communicate with fastboot (nonzero exit, output was [${output}])"
+    fi
+
+    if [[ "${output}" =~ "${serial}" ]]; then
+      found=true
+      break
+    fi
+
+    echo -n '.'
+    sleep 1
+  done
+  echo
+
+  if [[ -z "${found}" ]]; then
+    die "Couldn't find device with serial ${serial} in fastboot output"
+  else
+    echo "Found device ${serial}"
+  fi
 }
 
-flash_all_images() {
-    pushd ${PRODUCT_OUT}
-    erase_mmc
-    ${FASTBOOT_TOOL} flash mmc0 partition-table.img
-    ${FASTBOOT_TOOL} flash mmc0boot0 bl2.img
-    ${FASTBOOT_TOOL} flash mmc0boot1 u-boot-env.bin
-    ${FASTBOOT_TOOL} flash bootloaders fip.bin
-    ${FASTBOOT_TOOL} flash boot boot_${USERSPACE_ARCH}.img
-    ${FASTBOOT_TOOL} flash rootfs rootfs_${USERSPACE_ARCH}.img
-    popd
+function main {
+    local usage=$(cat <<EOF
+Usage: flash.sh [-d <files_dir>] [-s <serial>] [-r <detect_retries>]
+  -l                  load LK into ram for fastboot mode
+  -d <files_dir>      flashes files from <files_dir> (defaults to current dir)
+  -s <serial>         only flashes the board with the given fastboot serial number
+  -r <detect_retries> number of times to retry waiting for a device (defaults to 0)
+EOF
+)
+  local detect_retries=1    # -r <retry_count>
+  local serial_number       # -s <serial>
+  local files_dir           # -d <files_dir>
+  local lk_fastboot         # -l
+  local args=$(getopt hls:d:r: $*)
+  set -- $args
+
+  for i; do
+    case "$1" in
+      -l)  # lk
+        local lk_fastboot=true
+        shift 1
+        ;;
+
+      -s)  # serial
+        serial_number="$2"
+        shift 2
+        ;;
+
+      -d)  # files dir
+        files_dir="$2"
+        shift 2
+        ;;
+
+      -r)  # retry count
+        detect_retries="$2"
+        if ! [[ "${detect_retries}" -gt "0" ]]; then
+          die "Retry count must be a positive integer!"
+        fi
+        shift 2
+        ;;
+
+      --)  # end of args
+        shift
+        break
+        ;;
+
+      *)  # help
+        die "${usage}"
+        ;;
+
+    esac
+  done
+
+  if [[ ! -x "${FASTBOOT_CMD}" ]]; then
+    die "Couldn't find fastboot on your PATH -- did you install it?"
+  fi
+
+  if [[ -n "${serial_number}" ]]; then
+    FASTBOOT_CMD="${FASTBOOT_CMD} -s ${serial_number}"
+  fi
+
+  if [[ -z "${files_dir}" ]]; then
+    if [[ -n "${PRODUCT_OUT}" ]]; then
+      files_dir="${PRODUCT_OUT}"
+    else
+      files_dir="${SCRIPT_DIR}"
+    fi
+  fi
+
+  ensure_flash_files_present "${files_dir}"
+
+  if [[ -n "${lk_fastboot}" ]]; then
+    load_fastboot "${files_dir}"
+    exit 0
+  fi
+
+  if [[ -n "${serial_number}" ]]; then
+    detect_device_or_die "${detect_retries}" "${serial_number}"
+  fi
+
+  flash_partitions "${files_dir}"
+  try ${FASTBOOT_CMD} reboot
+
+  echo "Flash completed."
 }
 
-
-debug=0
-
-FLASHALL=0
-LKONLY=1
-mode=$FLASHALL
-
-while getopts "h?vbf" opt; do
-	case $opt in
-		h|\?)
-			show_help
-			exit 0
-			;;
-		v)
-			set -x
-			debug=1
-			;;
-		b)
-			mode=$BOOTONLY
-			;;
-		f)
-			mode=$LKONLY
-			;;
-	esac
-done
-
-load_fastboot
-
-case $mode in
-	$FLASHALL)
-		flash_all_images
-		sleep 1
-		${FASTBOOT_TOOL} reboot
-		exit 0
-		;;
-	$LKONLY)
-		flash_fastboot
-		sleep 1
-		exit 0
-		;;
-	*) echo "Unknown flash mode."
-esac
-
+main "$@"