#!/bin/bash
# Copyright 2017 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

### make a zedboot USB key

## usage: fx mkzedboot [options] <usb device>
##  -f            force writing to a non-usb target

set -e

source "$(cd "$(dirname "${BASH_SOURCE[0]}")" >/dev/null 2>&1 && pwd)"/lib/image_build_vars.sh || exit $?

if [[ "${FUCHSIA_ARCH}" != "x64" ]]; then
  echo >&2 mkzedboot is not supported for ${FUCHSIA_ARCH}
  exit 1
fi

force=false
if [[ "$1" == "-f" ]]; then
  shift
  force=true
fi

is_usb() {
  if ! ${force}; then
    fx-command-run list-usb-disks | grep "$1"
  fi
}

USB_DEVICE="$1"
if [[ -z "${USB_DEVICE}" ]]; then
  echo >&2 "device argument required"
  echo "USB disks:"
  fx-command-run list-usb-disks
  exit 1
fi
if ! is_usb "${USB_DEVICE}"; then
  echo >&2 "${USB_DEVICE} does not look like a USB device, use -f to force, or pick from below"
  echo "USB disks:"
  fx-command-run list-usb-disks
  exit 1
fi

if ! [[ -w "${USB_DEVICE}" ]] ; then
  echo >&2 "Changing ownership of ${USB_DEVICE} to ${USER}"
  sudo chown "${USER}" "${USB_DEVICE}"
fi

echo >&2 "Opening device..."
# We open the device and hold onto an fd for the duration of our modifications.
# This prevents automounting solutions from observing a final close and
# rescanning the partition table until we're all done making changes -
# particularly important on macOS where users would otherwise receive
# EAGAIN/EBUSY and so on.
open_device() {
  case "$(uname)" in
  Darwin)
    if ! diskutil quiet unmountDisk "${USB_DEVICE}"; then
      echo >&2 "Failed to unmount ${USB_DEVICE}, cannot continue"
      exit 1
    fi
    ;;
  esac
  exec 3>>"${USB_DEVICE}"
}
close_device() {
  echo >&2 "Closing device."
  exec 3>&-
}
open_device

# Destroy any existing GPT/MBR on the device and re-create
echo "Create new GPT partition table... "
"${FUCHSIA_BUILD_DIR}/tools/cgpt" create "${USB_DEVICE}"
"${FUCHSIA_BUILD_DIR}/tools/cgpt" boot -p "${USB_DEVICE}"
echo "done"

echo "Create new partitions... "
# ESP needs to be a FAT compatible size
if [[ ! -z "${IMAGE_ZEDBOOT_ESP}" ]]; then
  esp_size=$(((63*1024*1024)/512))
  esp_offset=2048
  "${FUCHSIA_BUILD_DIR}/tools/cgpt" add -s "${esp_size}" -t efi -b "${esp_offset}" -l esp "${USB_DEVICE}"
else
  reserved_size=2048
  reserved_offset=2048
  vboot_size=$(((64*1024*1024)/512))
  vboot_offset=$(($reserved_offset + $reserved_size))
  "${FUCHSIA_BUILD_DIR}/tools/cgpt" add -s "${reserved_size}" -t reserved -b "${reserved_offset}" -l reserved "${USB_DEVICE}"
  "${FUCHSIA_BUILD_DIR}/tools/cgpt" add -s "${vboot_size}" -t kernel -b "${vboot_offset}" -l zedboot "${USB_DEVICE}"
fi

if [[ ! -z "${IMAGE_ZIRCONR_SIGNEDZBI}" ]]; then
  "${FUCHSIA_BUILD_DIR}/tools/cgpt" add -i 2 -T 1 -S 1 -P 2 "${USB_DEVICE}"
fi
echo "done"

if [[ ! -z "${IMAGE_ZEDBOOT_ESP}" ]]; then
  echo "Writing zedboot for EFI"
  dd if="${FUCHSIA_BUILD_DIR}/${IMAGE_ZEDBOOT_ESP}" of="${USB_DEVICE}" seek=${esp_offset}
else
  echo "Writing zedboot for Cros"
  dd if="${FUCHSIA_BUILD_DIR}/${IMAGE_ZIRCONR_SIGNEDZBI}" of="${USB_DEVICE}" seek=${vboot_offset}
fi
echo "done"

close_device

case "$(uname)" in
  Linux)
    eject "${USB_DEVICE}"
    ;;
  Darwin)
    diskutil eject "${USB_DEVICE}"
    ;;
esac