#!/usr/bin/env bash
# Flash Raspberry Pi SD card images on your PC or Mac
# Stefan Scherer - scherer_stefan@icloud.com
#
# Linux initial version by Matt Williams - matt@matthewkwilliams.com
# MIT License

set -eo pipefail

error()
{
  echo "$1"
  exit "$2"
}

usage()
{
  cat << EOF
usage: $0 [options] [name-of-rpi.img]

Flash a local or remote Raspberry Pi SD card image.

OPTIONS:
  --help|-h       Show this message
  --bootconf|-C   Copy this config file to /boot/config.txt
  --config|-c     Copy this config file to /boot/device-init.yaml (or occidentalis.txt)
  --hostname|-n   Set hostname for this SD image
  --ssid|-s       Set WiFi SSID for this SD image
  --password|-p   Set WiFI password for this SD image
  --clusterlab|-l Start Cluster-Lab on boot: true or false
  --device|-d     Card device to flash to (e.g. /dev/sdb in Linux or /dev/disk2 in OSX)
  --force|-f      Force flash without security prompt (for automation)
  --userdata|-u   Copy this cloud-init file to /boot/user-data
  --metadata|-m   Copy this cloud-init file to /boot/meta-data

If no image is specified, the script will try to configure an existing
image. This is useful to try several configuration without the need to
rewrite the image every time.

For HypriotOS < v1.7.0:

The config file device-init.yaml should look like

hostname: black-pearl
wifi:
  interfaces:
    wlan0:
      ssid: "MyNetwork"
      password: "secret_password"

For HypriotOS v1.7.0 and higher:

The config file user-data config file is the cloud-init configuration.
See http://cloudinit.readthedocs.io/en/0.7.9/ for more details.
EOF
  exit 1
}

# translate long options to short
for arg
do
  delim=""
  case "${arg}" in
    --help) args="${args}-h ";;
    --verbose) args="${args}-v ";;
    --config) args="${args}-c ";;
    --hostname) args="${args}-n ";;
    --ssid) args="${args}-s ";;
    --password) args="${args}-p ";;
    --bootconf) args="${args}-C ";;
    --clusterlab) args="${args}-l ";;
    --device) args="${args}-d ";;
    --force) args="${args}-f ";;
    --userdata) args="${args}-u ";;
    --metadata) args="${args}-m ";;
    # pass through anything else
    *) [[ "${arg:0:1}" == "-" ]] || delim="\""
      args="${args}${delim}${arg}${delim} ";;
  esac
done
# reset the translated args
eval set -- "$args"
# now we can process with getopt
while getopts ":hc:n:s:p:C:l:d:fu:m:" opt; do
  case $opt in
    h)  usage ;;
    c)  CONFIG_FILE=$OPTARG ;;
    C)  BOOT_CONF=$OPTARG ;;
    n)  SD_HOSTNAME=$OPTARG ;;
    s)  WIFI_SSID=$OPTARG ;;
    p)  WIFI_PASSWORD=$OPTARG ;;
    l)  CLUSTERLAB=$OPTARG ;;
    d)  DEVICE=$OPTARG ;;
    f)  FORCE=1 ;;
    u)  USER_DATA=$OPTARG ;;
    m)  META_DATA=$OPTARG ;;
    \?) usage ;;
    :)
      echo "option -$OPTARG requires an argument"
      usage
    ;;
  esac
done
shift $((OPTIND -1))

beginswith() { case $2 in $1*) true;; *) false;; esac; }
endswith() { case $2 in *$1) true;; *) false;; esac; }

if [ $# -lt 1 ]; then
  usage
fi

if [[ "$1" == "--help" ]]; then
  usage
fi
image=$1

if [[ -z $image ]]; then
  CONFIGURE_ONLY=1
  echo "Configuration mode. No image will be written"
fi

filename=$(basename "${image}")
extension="${filename##*.}"
filename="${filename%.*}"

# Figure out our OS
if [[ -z "${OSTYPE}" ]]; then
  OSTYPE=$(uname -s)
fi

case "${OSTYPE}" in
  darwin*)
    alias grep="grep --color=never"
    size_opt="-f %z"
    bs_size=1m

    # Check that the system has all the needed binaries/requirements in place
    check_requirements() {
      ## NO-OP in Darwin
      true
    }

    # Try to identify the most likely device that the user will use to
    # write an image to.
    #
    # return _RET: the name of the device to use
    autodetect_device() {
      set +e
      _RET=/dev/$(diskutil list | grep --color=never FDisk_partition_scheme | awk 'NF>1{print $NF}')

      if [ "${_RET}" == "" ] || [ "${_RET}" == "/dev/" ]; then
        echo "No SD card found. Please insert SD card, I'll wait for it..."
        while [ "${_RET}" == "" ] || [ "${_RET}" == "/dev/" ]; do
          sleep 1
          _RET=/dev/$(diskutil list | grep --color=never FDisk_partition_scheme | awk 'NF>1{print $NF}')
        done
      fi
      set -e
    }

    # Show in the standard output the devices that are a likely
    # destination for the tool to write an image to.
    show_devices() {
      diskutil list | grep FDisk_partition_scheme | awk 'NF>1{print $NF}'
    }

    # Check that the target device can be written. It will return 0 in
    # this case and 1 if it is not writable
    #
    # @param arg1 device name to check
    check_device_is_writable() {
      disk=$1
      if [[ "$disk" == *.dmg ]]; then
        # CI helper
        _RET=1
        return
      fi
      readonlymedia=$(diskutil info "$disk" | grep "Read-Only Media" | awk 'NF>1{print $NF}')
      if [[ $readonlymedia == "No" ]] ; then
        _RET=1
      else
        _RET=0
      fi
    }

    # Convert the device name into a raw device name that is suitable for
    # use by dd
    #
    # @param arg1 device name
    # @return _RET the raw device name
    get_raw_device_filename() {
      _RET="${1//\/dev\///dev/r}"
    }

    # Get the directory where the boot volume will be mounted.
    #
    # @param arg1 the name of the device holding the volume to be mounted
    # @return _RET mount point name
    get_boot_mount_point() {
      _RET=$(df | grep --color=never "${1}s1" | /usr/bin/sed 's,.*/Volumes,/Volumes,')
    }

    # Wait for the new created disk to be available
    #
    # @param arg1 device name to check
    wait_for_disk() {
      # helper for macOS CI
      rawdisk="$1"
      if [[ "$rawdisk" == *.dmg ]]; then
        mv "$rawdisk" "${rawdisk}.readonly.dmg"
        hdiutil convert "${rawdisk}.readonly.dmg" -format UDRW -o "$rawdisk"
        rm -f "${rawdisk}.readonly.dmg"
        disk=$(hdiutil attach "$rawdisk" | grep FAT | sed 's/s1 .*$//')
        echo mounted FAT partition to "$disk"
        if [ "$disk" == "" ]; then
          echo Failed attaching "$rawdisk"
          exit 5
        fi
      fi

      set +e
      find_boot_dev_name "$rawdisk"
      boot=$_RET
      if [ "${boot}" == "" ]; then
        COUNTER=0
        while [ $COUNTER -lt 5 ]; do
          sleep 1
          find_boot_dev_name "$rawdisk"
          boot=$_RET
          if [ "${boot}" != "" ]; then
            break
          fi
          (( COUNTER=COUNTER+1 ))
        done
      fi
      set -e
    }

    # Find the device name of the boot partition
    #
    # @param arg1 the disk name containing the partition
    find_boot_dev_name() {
      _RET=$(df | grep --color=never "${disk}s1" | /usr/bin/sed 's,.*/Volumes,/Volumes,')
    }

    # Unmount a disk
    #
    # @param arg1 the disk to unmount
    umount_disk() {
      disk=$1
      if [[ "$disk" == *.dmg ]]; then
        return
      fi
      set +e
      diskutil unmountDisk "${disk}s1"
      set -e
      diskutil unmountDisk "${disk}"
    }

    detach() {
      hdiutil detach "${1}"
    }
    # Mount the boot disk in the specified mount point
    #
    # @param arg1 the device to mount. The boot partition will be found automatically
    # @param arg2 mount point
    mount_boot_disk() {
      # NO-OP: Darwin will mount the boot partition automatically as soon
      # as the new disk is detected
      true
    }

    prepare_raw_disk() {
      _RET=$1
    }

    cleanup() {
      true
    }

    sudo_prompt() {
      # Do not use sudo -v otherwise Travis CI will hang.
      true
    }

    play_ok() {
      afplay /System/Library/Sounds/Bottle.aiff
    }

    play_warn() {
      afplay /System/Library/Sounds/Basso.aiff
    }

    sed_i() {
      sed -i "" "$@"
    }
    ;;
  Linux|linux|linux-gnu*)
    size_opt="-c %s"
    bs_size=1M

    # Check that the system has all the needed binaries/requirements in place
    check_requirements() {
      if ! sudo sh -c 'command -v hdparm' > /dev/null; then
        echo "No 'hdparm' command found; please install hdparm by running:"
        echo "sudo [apt-get|yum|something-else] install hdparm"
        exit 1
      fi
    }

    # Try to identify the most likely device that the user will use to
    # write an image to.
    #
    # @return _RET the name of the device to use
    autodetect_device() {
      _RET=$(lsblk -n -o NAME -d | grep mmcblk)
    }

    # Show in the standard output the devices that are a likely
    # destination for the tool to write an image to.
    show_devices() {
      if [[ -x $(command -v lsblk) ]]; then
        lsblk --output NAME,SIZE,TYPE,MOUNTPOINT
      else
        df -h
      fi
    }

    # Convert a image file into a destination disk
    #
    # @param arg1 the destination image file
    # @return _RET the disk that represents the image
    prepare_raw_disk() {
      if [[ "$1" == *.img ]]; then
        error "Raw files not supported under Linux yet" 2
      fi
      _RET=$1
    }

    # Convert the device name into a raw device name that is suitable for
    # use by dd
    #
    # @param arg1 device name
    # @return _RET the raw device name
    get_raw_device_filename() {
      _RET="${1}"
    }

    # Check that the target device can be written. It will return 0 in
    # this case and 1 if it is not writable
    #
    # @param arg1: device name to check
    check_device_is_writable() {
      disk=$1
      if [[ "$disk" == "loo" ]]; then
        # CI helper
        _RET=1
        return
      fi

      if sudo hdparm -r "$disk" | grep -q off; then
        _RET=1
      else
        _RET=0
      fi
    }

    # Get the directory where the boot volume will be mounted
    #
    # @param arg1 the name of the device holding the volume to be mounted
    # @return _RET: mount point name
    get_boot_mount_point() {
      _RET=/tmp/"$(id -u)"/mnt.$$
      mkdir -p "${_RET}"
    }

    # Wait for the new created disk to be available
    #
    # @param arg1 device name to check
    wait_for_disk() {
      echo "Waiting for device $1"
      udevadm settle
      sudo hdparm -z "$1"
    }

    # Find the device name of the boot partition
    #
    # @param arg1 the disk name containing the partition
    find_boot_dev_name() {
      if beginswith /dev/mmcblk "${1}" ;then
        _RET="${1}p1"
      elif beginswith /dev/loop "${1}" ;then
        _RET="${1}p1"
      else
        _RET="${1}1"
      fi
    }


    # Unmount a disk
    #
    # @param arg1 the disk to unmount
    umount_disk() {
      for i in $(df |grep "$1" | awk '{print $1}')
      do
        sudo umount "$i"
      done
    }

    detach() {
      umount_disk "$1"
    }

    # Mount the boot disk in the specified mount point
    #
    # @param arg1 the device to mount. The boot partition will be found automatically
    # @param arg2 mount point
    mount_boot_disk() {
      local disk=$1
      local mount_point=$2
      local dev

      find_boot_dev_name "${disk}"
      dev=$_RET

      sudo mount -o uid="$(id -u)",gid="$(id -g)" "${dev}" "${mount_point}"
      ls -la "${mount_point}"
    }

    cleanup() {
      rmdir "$1"
    }

    sudo_prompt() {
      # this sudo here is used for a login without pv's progress bar
      # hiding the password prompt
      sudo -v
    }

    play_ok() {
      true
    }

    play_warn() {
      true
    }

    sed_i() {
      sed -i "$@"
    }
    ;;
  *)
    echo Unknown OS: "${OSTYPE}"
    exit 11
    ;;
esac

if endswith Microsoft "$(uname -r)"; then
  echo This script does not work in WSL.
  exit 11
fi

check_requirements

if [ ! -z "${USER_DATA}" ]; then
  if [ ! -f "${USER_DATA}" ]; then
    echo "Cloud-init file ${USER_DATA} not found!"
    exit 10
  fi
fi

if [ ! -z "${META_DATA}" ]; then
  if [ ! -f "${META_DATA}" ]; then
    echo "Cloud-init file ${META_DATA} not found!"
    exit 10
  fi
fi

if [ ! -z "${BOOT_CONF}" ]; then
  if [ ! -f "${BOOT_CONF}" ]; then
    echo "File ${BOOT_CONF} not found!"
    exit 10
  fi
fi

if [ ! -z "${CONFIG_FILE}" ]; then
  if [ ! -f "${CONFIG_FILE}" ]; then
    echo "File ${CONFIG_FILE} not found!"
    exit 10
  fi
fi

if [[ -z $CONFIGURE_ONLY ]] ; then
  if [ -f "/tmp/${filename}" ]; then
    image=/tmp/${filename}
    echo "Using cached image ${image}"
  elif [ -f "/tmp/${filename}.img" ]; then
    image=/tmp/${filename}.img
    echo "Using cached image ${image}"
  else
    if beginswith http:// "${image}" || beginswith https:// "${image}"; then
      command -v curl 2>/dev/null || error "Error: curl not found. Aborting" 1
      echo "Downloading ${image} ..."
      curl -L --fail -o "/tmp/image.img.${extension}" "${image}"
      image=/tmp/image.img.${extension}
    fi

    if beginswith s3:// "${image}" ;then
      command -v aws 2>/dev/null || error "Error: aws not found. Aborting" 1
      echo "Downloading ${image} ..."
      aws s3 cp "${image}" "/tmp/image.img.${extension}"
      image=/tmp/image.img.${extension}
    fi

    if [ ! -f "${image}" ]; then
      echo "File ${image} not found."
      exit 10
    fi

    if [[ "$(file "${image}")" == *"Zip archive"* ]]; then
      command -v unzip 2>/dev/null || error "Error: unzip not found. Aborting" 1
      echo "Uncompressing ${image} ..."
      unzip -o "${image}" -d /tmp
      image=$(unzip -l "${image}" | grep -v Archive: | grep img | awk 'NF>1{print $NF}')
      image="/tmp/${image}"
      echo "Use ${image}"
    fi

    if [[ "$(file "${image}")" == *"gzip compressed data"* ]]; then
      echo "Uncompressing ${image} ..."
      gzip -d "${image}" -c >/tmp/image.img
      image=/tmp/image.img
      echo "Use ${image}"
    fi

    if [[ "$(file "${image}")" == *"xz compressed data"* ]]; then
      command -v xz 2>/dev/null || error "Error: unzip not found. Aborting" 1
      echo "Uncompressing ${image} ..."
      xz -d "${image}" -c >/tmp/image.img
      image=/tmp/image.img
      echo "Use ${image}"
    fi
  fi
fi

while true; do
  disk="$DEVICE"
  if [[ -z "${disk}" ]]; then
    autodetect_device
    disk="$_RET"

    if [[ -z "${disk}" ]]; then
      show_devices
      # shellcheck disable=SC2162
      read -p "Please pick your device: "
      disk="${REPLY}"
      [[ ${disk} != /dev/* ]] && disk="/dev/${disk}"
    fi
  fi
  if [[ -z "${FORCE}" ]]; then
    while true; do
      echo ""
      read -rp "Is ${disk} correct? " yn
      case $yn in
        [Yy]* ) break;;
        [Nn]* ) exit;;
        * ) echo "Please answer yes or no.";;
      esac
    done
  fi

  prepare_raw_disk "${disk}"
  disk=$_RET

  check_device_is_writable "${disk}"
  writable=$_RET

  echo "Unmounting ${disk} ..."
  umount_disk "${disk}"

  if [ "$writable" == "1" ]; then
    break
  else
    play_warn
    echo "The SD card is write protected. Please eject, remove protection and insert again."
  fi
done

if [[ -z $CONFIGURE_ONLY ]] ; then
  get_raw_device_filename "$disk"
  rawdisk=$_RET

  echo "Flashing ${image} to ${rawdisk} ..."
  if [[ -x $(command -v pv) ]]; then
    sudo_prompt
    size=$(/usr/bin/stat "$size_opt" "${image}")
    pv -s "${size}" < "${image}" | sudo dd bs=$bs_size "of=${rawdisk}"
  else
    echo "No 'pv' command found, so no progress available."
    echo "Press CTRL+T if you want to see the current info of dd command."
    sudo dd bs=1M "if=${image}" "of=${rawdisk}"
  fi

  wait_for_disk "${disk}"
fi

echo "Mounting Disk"
get_boot_mount_point "${disk}"
boot="$_RET"

echo "Mounting ${disk} to customize..."
mount_boot_disk "${disk}" "${boot}"

 if [ -f "${CONFIG_FILE}" ]; then
    if [[ "${CONFIG_FILE}" == *"occi"* ]]; then
      echo "Copying ${CONFIG_FILE} to ${boot}/occidentalis.txt ..."
      cp "${CONFIG_FILE}"  "${boot}/occidentalis.txt"
    else
      echo "Copying ${CONFIG_FILE} to ${boot}/device-init.yaml ..."
      cp "${CONFIG_FILE}"  "${boot}/device-init.yaml"
    fi
  fi

  if [[ -f "${BOOT_CONF}" ]]; then
    echo "Copying ${BOOT_CONF} to ${boot}/config.txt ..."
    cp "${BOOT_CONF}" "${boot}/config.txt"
  fi

  if [ -f "${USER_DATA}" ]; then
    echo "Copying cloud-init ${USER_DATA} to ${boot}/user-data ..."
    cp "${USER_DATA}" "${boot}/user-data"
  fi

  if [ -f "${META_DATA}" ]; then
    echo "Copying cloud-init ${META_DATA} to ${boot}/meta-data ..."
    cp "${META_DATA}" "${boot}/meta-data"
  fi


if [ -f "${boot}/device-init.yaml" ]; then
  echo "Setting device-init"
  if [ ! -z "${SD_HOSTNAME}" ]; then
    echo "  Set hostname=${SD_HOSTNAME}"
    sed_i -e "s/.*hostname:.*\$/hostname: ${SD_HOSTNAME}/" "${boot}/device-init.yaml"
  fi
  if [ ! -z "${WIFI_SSID}" ]; then
    echo "  Set wlan0/ssid=${WIFI_SSID}"
    sed_i -e "s/.*wlan0:.*\$/    wlan0:/" "${boot}/device-init.yaml"
    sed_i -e "s/.*ssid:.*\$/      ssid: \"${WIFI_SSID}\"/" "${boot}/device-init.yaml"
  fi
  if [ ! -z "${WIFI_PASSWORD}" ]; then
    echo "  Set wlan0/password=${WIFI_PASSWORD}"
    sed_i -e "s/.*wlan0:.*\$/    wlan0:/" "${boot}/device-init.yaml"
    sed_i -e "s/.*password:.*\$/      password: \"${WIFI_PASSWORD}\"/" "${boot}/device-init.yaml"
  fi
  if [ ! -z "${CLUSTERLAB}" ]; then
    echo "  Set Cluster-Lab/run_on_boot=${CLUSTERLAB}"
    sed_i -e "s/.*run_on_boot.*\$/    run_on_boot: \"${CLUSTERLAB}\"/" "${boot}/device-init.yaml"
  fi
fi

 # cloud-init
  if [ -f "${boot}/user-data" ]; then
    if [ ! -z "${SD_HOSTNAME}" ]; then
      echo "Set hostname=${SD_HOSTNAME}"
      sed_i -e "s/.*hostname:.*\$/hostname: ${SD_HOSTNAME}/" "${boot}/user-data"
    fi

    if [ ! -f "${boot}/meta-data" ]; then
      echo "Creating empty meta-data"
      touch "${boot}/meta-data"
    fi
  fi
# legacy: /boot/occidentalis.txt of old Hector release
if [ -f "${boot}/occidentalis.txt" ]; then
  echo "Setting Occidentalis"
  if [ ! -z "${SD_HOSTNAME}" ]; then
    echo "  Set hostname=${SD_HOSTNAME}"
    sed_i -e "s/.*hostname.*=.*\$/hostname=${SD_HOSTNAME}/" "${boot}/occidentalis.txt"
  fi
  if [ ! -z "${WIFI_SSID}" ]; then
    echo "Set wifi_ssid=${WIFI_SSID}"
    sed_i -e "s/.*wifi_ssid.*=.*\$/wifi_ssid=${WIFI_SSID}/" "${boot}/occidentalis.txt"
  fi
  if [ ! -z "${WIFI_PASSWORD}" ]; then
    echo "Set wifi_password=${WIFI_PASSWORD}"
    sed_i -e "s/.*wifi_password.*=.*\$/wifi_password=${WIFI_PASSWORD}/" "${boot}/occidentalis.txt"
  fi
fi

echo "Unmounting ${disk} ..."
sleep 1

set +e
detach "${disk}"

# shellcheck disable=SC2181
if [ $? -eq 0 ]; then
  cleanup "${boot}"
  play_ok
  echo "Finished."
else
  play_warn
  echo "Something went wrong."
fi