#!/bin/sh

SCRIPT_NAME=$(basename "$0")
UPGRADE_RES="N"

SANITIZE_TARGET="n"
REBOOT_TARGET="n"
PRE_REBOOT_TARGET="n"

SSH_COMMON_OPTS=""
SSH_UNSAFE_OPTS="-o UserKnownHostsFile=/dev/null  -o StrictHostKeyChecking=no"

# Color definitions
GREEN="\\033[1;32m"
NORMAL="\\033[0;39m"
RED="\\033[1;31m"
PINK="\\033[1;35m"
BLUE="\\033[1;34m"
WHITE="\\033[0;02m"
CLEARWHITE="\\033[1;08m"
YELLOW="\\033[1;33m"
CYAN="\\033[1;36m"

error()
{
  echo "Error: $1" >&2
  exit 1
}

usage()
{
    echo ""
    echo "Usage: '$SCRIPT_NAME' [options]"
    echo "Upgrade tool"
    echo ""
    echo "Mandatory arguments are -f, -i, -u and -p"
    echo "-h : show help"
    echo "-f <path> : path of the upgrade file"
    echo "-u <username> : admin username"
    echo "-p <password> : admin password"
    echo "-i <ipaddress> : IP address of the Network card to be upgraded"
    echo "-r : reboot the card after the upgrade"
    echo "-b : reboot the card before the upgrade"
    echo "-s : sanitize the card after the upgrade (card will reboot)"
    echo "-k : allow connection to host with unknown public key (UNSAFE)"
    echo ""
}

while getopts "p:i:f:u:rbshk" option; do
    case "$option" in
        f) IMAGE_PATH=$OPTARG ;;
        u) USERNAME=$OPTARG ;;
        p) USERPASSWORD=$OPTARG ;;
        i) TARGET_IP=$OPTARG ;;
        r) REBOOT_TARGET="y" ;;
        b) PRE_REBOOT_TARGET="y" ;;
        s) SANITIZE_TARGET="y" ;;
        k) SSH_COMMON_OPTS="${SSH_UNSAFE_OPTS}" ;;
        :) error "missing argument";;
        h) usage;exit 0 ;;
        \?) error "unrecognized option or missing arg";;
    esac
done

if [ -z "${IMAGE_PATH}" ] || [ -z "${TARGET_IP}" ] || [ -z "${USERNAME}" ]|| [ -z "${USERPASSWORD}" ]; then
  error "missing argument, -f <path> -i <ipaddress> -u <username> -p <password> mandatory"
fi

##########################################################
# Functions
##########################################################

# Print information on screen with color
writeInfo()
{
  printf "%s \033[0K\r" "${1}"
}

# Sent upgrade file by scp to the target
sendUpgradeFile()
{
  writeInfo "Transfering (${IMAGE_PATH}) with SCP to [${TARGET_IP}]"
  # With newer version of openssh, we have to use -O option for scp to not do
  # default switch to sftp protocol (not supported on the board) but use the 'legacy' scp protocol
  timeout 1m sshpass -p ${USERPASSWORD} scp -O ${SSH_COMMON_OPTS} ${IMAGE_PATH} ${USERNAME}@${TARGET_IP}:/tmp/download/enmcUpdate.tar 2> /dev/null
  if [ "$?" -ne "0" ]; then
  	  # Try to upload the file with the sftp protocol instead of the legacy scp one.
	  timeout 1m sshpass -p ${USERPASSWORD} scp ${SSH_COMMON_OPTS} ${IMAGE_PATH} ${USERNAME}@${TARGET_IP}:/tmp/download/enmcUpdate.tar 2> /dev/null
    if [ "$?" -ne "0" ]; then
	    error "Transfer failed of [${IMAGE_PATH}] to [${TARGET_IP}]"
	  else
	    writeInfo "Transfer done."
	  fi
  else
    writeInfo "Transfer done."
  fi
}

# Get status during upgrade process
keepWatchingUpgradeProgress()
{
  CONTINUE="Y"
  writeInfo "Checking upgrade status ..."
  sleep 5

  # Manage timer to check process
  TIME_OUT_DURATION=90
  STARTING_UPGRADE=$(date  +%s)
  TIME_OUT=$((STARTING_UPGRADE+TIME_OUT_DURATION))

  while [ "${CONTINUE}" = "Y" ]
  do
    # Get status (step, progression and error if any occurs) from target through ssh
    UPGRADE_STEP=$(sshpass -p ${USERPASSWORD} ssh ${SSH_COMMON_OPTS} ${USERNAME}@${TARGET_IP} 'firmware_update_status step' 2> /dev/null)
    UPGRADE_PROGRESSION=$(sshpass -p ${USERPASSWORD} ssh ${SSH_COMMON_OPTS} ${USERNAME}@${TARGET_IP} 'firmware_update_status progression' 2> /dev/null)
    UPGRADE_ERROR=$(sshpass -p ${USERPASSWORD} ssh ${SSH_COMMON_OPTS} ${USERNAME}@${TARGET_IP} 'firmware_update_status error' 2> /dev/null)

    TIME_NOW=$(date  +%s)

    # Show status and check time out
    LOG_MESSAGE=""
    if [ -n "${UPGRADE_STEP}" ]; then
      LOG_MESSAGE="${UPGRADE_STEP}"
      if [ -n "${UPGRADE_PROGRESSION}" ]; then
        LOG_MESSAGE="${LOG_MESSAGE}"" - ""${UPGRADE_PROGRESSION}%%"
      fi
    fi
    if [ "${TIME_NOW}" -gt "${TIME_OUT}" ]; then
      echo ""
      echo "Timeout exceeded: the upgrade process is not responding after ${TIME_OUT_DURATION} seconds."
      CONTINUE="N"
      UPGRADE_RES="N"
      writeResultFailed "upgrade" "No response after timeout"
    else
      LOG_MESSAGE="Upgrade ..."
    fi

    writeInfo "${LOG_MESSAGE}"

    if [ -n "${UPGRADE_ERROR}" ]; then
      CONTINUE="N"
      UPGRADE_RES="N"
      writeResultFailed "Upgrade Progress" "${UPGRADE_ERROR}"
    else

      UPGRADE_DONE=$(echo "${UPGRADE_STEP}" | grep "upgrade done")

      if [ -n "$UPGRADE_DONE" ] && [ "${UPGRADE_PROGRESSION}" -ge 100 ]; then
        UPGRADE_RES="Y"
        CONTINUE="N"
      fi
    fi
  done
}

preRebootCard()
{
  if [ "${PRE_REBOOT_TARGET:-n}" = "n" ]; then
    writeInfo "Pre reboot was not requested."
  else
    
    # Reboot the card
    writeInfo "Rebooting..."
    TIME_OUT_DURATION=300
    STARTING_REBOOT=$(date  +%s)
    TIME_OUT=$((STARTING_REBOOT+TIME_OUT_DURATION))
    if sshpass -p ${USERPASSWORD} ssh ${SSH_COMMON_OPTS} -o LogLevel=QUIET ${USERNAME}@${TARGET_IP} 'reboot --withoutconfirmation'; then
      echo ""
      printf "Waiting for the card to reboot"
      sleep 20
    else
      echo "Error when rebooting the card"
      exit 1
    fi
    
    # Wait for the card to reboot or timeout
    CONTINUE="Y"
    while [ "${CONTINUE}" = "Y" ]
    do
      TIME_NOW=$(date  +%s)
      if [ "${TIME_NOW}" -gt "${TIME_OUT}" ]; then
        echo ""
        echo "Timeout exceeded: the IP address is not available after ${TIME_OUT_DURATION} seconds."
        exit 1
      fi
      if sshpass -p ${USERPASSWORD} ssh ${SSH_COMMON_OPTS} -o ConnectTimeout=2 -o LogLevel=QUIET ${USERNAME}@${TARGET_IP} 'logout'; then
        CONTINUE="N"
        echo ""
        echo "The card has finished rebooting"
      else
        printf "."
      fi
      sleep 1
    done

  fi
  echo ""
}

setupPostRebootCard()
{
  if [ "$REBOOT_TARGET" = "y" ]; then
    sshpass -p ${USERPASSWORD} ssh ${SSH_COMMON_OPTS} ${USERNAME}@${TARGET_IP} 'reboot --afterupdate' 2> /dev/null
  fi
}

rebootCard()
{
  if [ -z "$REBOOT_TARGET" ] || [ "$REBOOT_TARGET" = "n" ]; then
    writeInfo "Reboot was not requested."
  else
    TIME_OUT_DURATION=300
    STARTING_REBOOT=$(date  +%s)
    TIME_OUT=$((STARTING_REBOOT+TIME_OUT_DURATION))
    writeInfo "Rebooting..."
  fi
  echo ""
}

sanitize()
{
  if [ -z "${SANITIZE_TARGET}" ] || [ "${SANITIZE_TARGET}" = "n" ]; then
    writeInfo "Sanitize was not requested."
  else
    writeInfo "Sanitizing..."
    sshpass -p ${USERPASSWORD} ssh ${SSH_COMMON_OPTS} ${USERNAME}@${TARGET_IP} 'sanitize --withoutconfirmation' 2> /dev/null
    writeResult
  fi
  echo ""
}

# Arg 1 step, arg2 error comment
writeResultFailed()
{
  printf "%bUpgrade:%b Failed during %s, error: %s              %b\n" "${CYAN}" "${RED}" "${1}" "${2}" "${NORMAL}"
  error "Upgrade has failed during ${1}, error: ${2}"
}

writeResultOK()
{
  printf "%bUpgrade:%b OK%b\n" "${CYAN}" "${GREEN}" "${NORMAL}"
  exit 0
}

writeHeader()
{
  printf "\n%bSTARTING UPGRADE FROM: [%s] to [%s]%b\n" "${CYAN}" "${IMAGE_PATH}" "${TARGET_IP}" "${NORMAL}"
}

writeResult()
{
  writeInfo "" # cleanup the line
  if [ "${UPGRADE_RES}" = "Y" ];
  then writeResultOK
  else writeResultFailed
  fi
}

checkConnection()
{
  sshpass -p ${USERPASSWORD} ssh ${SSH_COMMON_OPTS} -o ConnectTimeout=2  -o LogLevel=QUIET ${USERNAME}@${TARGET_IP} 'logout';
  # from sshpass manual: 6 -> Host public key is unknown. sshpass exits without confirming the new key
  if [ $? -eq 6 ]; then
    echo "${RED}Target (${TARGET_IP}) host public key is unknown.${NORMAL}"
    echo "Add the key to known hosts or run the script with option -k to accept unknown key."
    exit 6
  fi
}

checkConnection
writeHeader
preRebootCard
setupPostRebootCard
sendUpgradeFile
keepWatchingUpgradeProgress
sanitize
rebootCard
writeResult
