#!/bin/bash

#
# File .......... swuk_upload_boot
# Author ........ Steve Haywood
# Website ....... http://www.spacewire.co.uk
# Project ....... Common (SpaceWire UK Tutorial)
# Date .......... 05 Mar 2026
# Version ....... 2.0
# Description ...
#   Upload boot files (BOOT.BIN, boot.src & image.ub) to a remote system.
# Process involves backing up the boot files on the remote system, uploading the
# new ones and rebooting the system.
#

# Strict
#set -euo pipefail


################################################################################
# Check existence of boot files
# Arguments ...... $1 ... Str: Files directory
# Return ......... None
# Exit code ...... None
# Shared (In) .... c_bfiles
# Shared (Out) ... None
check_boot_files () {
  # Declare variables
  local -r dir=${1}  # Str: Files directory
  local    bfile     # Str: Current file from the directory

  # Iterate through files
  for bfile in "${c_bfiles[@]}"
  do
    echo -n " - "
    [ -f ${dir}/${bfile} ] && echo -en "${c_green}" || echo -en "${c_red}";
    echo -en "${bfile}${c_normal}"
  done
  echo
}


################################################################################
# Main function.
# Arguments ...... ${@}
# Return ......... None
# Exit code ...... Status (0=success, 1=failure)
# Shared (In) .... None
# Shared (Out) ... None
main()
{
  # Declare terminal colour constants
  local -r  c_red='\033[0;31m'
  local -r  c_green='\033[0;32m'
  local -r  c_normal='\033[0m'

  # Declare argument constants
  local -r  c_opthelp="--help"  # Str: Help option name
  local -Ar c_options=(         # ARR: Options & associated help information
    [${c_opthelp}]="Display this help and exit."
  )
  local -ar c_optorder=(        # Arr: Help options display order
    ${c_opthelp}
  )
  local -ar argv=(${@})         # Arr: Get argument values (space-separated) into array

  local -r  hosts="${HOME}/.ssh/known_hosts"             # Str: Known hosts file
  local -r  c_ipre='^(0*(1?[0-9]{1,2}|2([0-4][0-9]|5[0-5]))\.){3}0*(1?[0-9]{1,2}|2([0-4][0-9]|5[0-5]))$'
                                                         # Str: IP validation regular expression
  local -r  c_local="/tftpboot"                          # Str: Location of most recent OS files
  local -r  c_archive="${swuk_user}/petalinux"           # Str: Location of user OS files
  local -r  c_remote="/media/sd-mmcblk0p1"               # Str: Remote SD-Card location
  local -r  c_bfiles=("BOOT.BIN" "boot.scr" "image.ub")  # Str: Boot files
  local -ri c_timedout=30                                # Int: Timeout limit (wait, ping, ssh)

  # Declare variables
  local     arg                      # Str: Current argument from argv array
  local     option                   # Str: Current option from c_optorder array
  local     ip=""                    # Str: IP Address (x.x.x.x)
  local -i  selcnt=0                 # Int: Selection count
  local -A  table                    # ARR: Selection table
  local -a  dirs                     # Arr: List of archive directories
  local     dir                      # Str: Current archive directory
  local     sel                      # Str: User selection input
  local     fpath                    # Str: Boot files location (path)
  local     bfile                    # Str: Boot file from array
  local -i  timeout                  # Int: Timeout for reboot (wait, ping, ssh)
  local     username="root"          # Str: PetaLinux username
  local     password="root"          # Str: PetaLinux password
  local     sudo="sudo"              # Str: Requires sudo for reboot

  # Display help information
  if [[ " ${argv[*]} " =~ " ${c_opthelp} " ]]; then
    echo "Usage: $(basename ${0}) IP-ADDRESS... [OPTION]..."
    echo "Upload the boot files ${c_bfiles[0]}, ${c_bfiles[1]} & ${c_bfiles[2]} to the remote system located at IP-ADDRESS."
    echo
    for option in ${c_optorder[@]}
    do
      echo "      ${option} $(printf ' %.0s' {1..12} | head -c $((12-${#option}))) ${c_options[${option}]}"
    done
    echo
    exit 0
  fi

  # Get & check the arguments
  for arg in ${argv[@]}; do
    if [[ ${arg:0:2} == "--" ]]; then  # Option
      [[ ! -v c_options[${arg}] ]] &&
        echo "Option (${arg}) is not recognised!" >&2 &&
        exit 1
    elif [[ -z ${ip} ]]; then  # IP address
      ip=${arg}
    else
      echo "Unexpected argument (${arg}) found!" >&2
      exit 1
    fi
  done

  # Further check the arguments
  [[ -z ${ip} ]] &&
    echo "No IP address specified!" >&2 &&
    exit 1

  # Check the IP ADDRESS
  [[ ! ${ip} =~ ${c_ipre} ]] &&
    echo "The IP address (${ip}) is incorrectly specified!" >&2 &&
    exit 1

  # Check directory at local location
  echo -e "\nBoot files @ ${c_local} :-\n"
  if [ -d "${c_local}" ]; then
    # Add local to selection list
    table[${selcnt}]=${c_local}
    # Check existence of boot files
    echo -n "${selcnt})"
    check_boot_files ${c_local}
    # Increase selection count
    selcnt=$((selcnt+1))
  else
    echo "Directory does not exist!"
  fi

  # Check directory at archive location
  echo -e "\nBoot files @ ${c_archive} :-\n"
  if [ -d "${c_archive}" ]; then
    # Get list of directories at archive location
    readarray -t dirs < <(find "${c_archive}" -mindepth 1 -maxdepth 1 -type d -printf '%P\n')
    # Check existence of subdirectories
    if [ ${#dirs[@]} -ne 0 ]; then
      # Sort list of directories
      dirs=($(IFS=$'\n'; echo "${dirs[*]}" | sort))
      # List what directories are present
      for dir in "${dirs[@]}"; do
        # Add archive subdirectory to selection list
        table[${selcnt}]="${c_archive}/${dir}"
        # Check existence of boot files
        echo -n "${selcnt}) ${dir}"
        check_boot_files "${c_archive}/${dir}"
        # Increase selection count
        selcnt=$((selcnt+1))
      done
    else
      echo "Directory does not contain any subdirectories!"
    fi
  else
    echo "Directory does not exist!"
  fi

  # Obtain user input, if good attempt upload
  echo -e "\nq) Quit\n"
  echo -n "Select boot files for upload or quit"
  while true; do
    read -p " : " sel
    if [[ ("${sel}" == "q") || ("${sel}" == "Q") ]]; then
      echo "Quitting without selection!"
      exit 0
    elif ! [[ "${sel}" =~ ^[0-9]+$ ]] ; then
      # User input is not a number
      echo -n "Selection is not a number, try again"
    elif [ ${sel} -lt 0 ] || [ ${sel} -ge ${selcnt} ]; then
      # User input is out of range
      echo -n "Selection is out of range, try again"
    else
      break
    fi
  done

  # Get path from table
  fpath=${table[${sel}]}

  # Iterate through list of boot files to check for existence
  for bfile in "${c_bfiles[@]}"; do
    [ ! -f ${fpath}/${bfile} ] &&
      echo -e "Boot file ($bfile) not found!" >&2 &&
      exit 1
  done

  # Get user credentials
  echo
  echo "Enter user credentials for ${ip} :-"
  read -e -p "Username: " -i "${username}" username
  password=${username}  # Have a guess at the password
  read -e -p "Password: " -i "${password}" password

  # Establish SSH connection
  swuk_ssh ${ip} > /dev/null 2>&1
  [ ${?} -ne 0 ] && echo "Failed to setup ssh!" >&2 && exit 1
  # Backup boot files
  sshpass -p ${password} ssh -t ${username}@${ip} "cd ${c_remote}; for file in ${c_bfiles[0]} ${c_bfiles[1]} ${c_bfiles[2]}; do cp \$file \$file.bak; done" > /dev/null 2>&1
  [ ${?} -ne 0 ] && echo "Failed to backup BOOT files!" >&2 && exit 1

  # Upload new boot files
  sshpass -p ${password} scp ${fpath}/{${c_bfiles[0]},${c_bfiles[1]},${c_bfiles[2]}} ${username}@${ip}:${c_remote} > /dev/null 2>&1
  [ ${?} -ne 0 ] && echo "Failed to upload BOOT files!" >&2 && exit 1

  # Reboot system
  [[ ${username} == "root" ]] && sudo=""
  sshpass -p ${password} ssh -t ${username}@${ip} "${sudo} /sbin/reboot" 2> /dev/null
  [ ${?} -ne 0 ] && echo "Failed to reboot PetaLinux!" >&2 && exit 1

  # Allow time for system to reboot
  echo
  tput civis
  for ((timeout=0; timeout<=${c_timedout}; timeout++)); do
    echo -ne "Remote reboot (wait): $timeout\033[0K / ${c_timedout}s\r"
    sleep 1
  done
  tput cnorm
  echo

  tput civis
  # Check IP-ADDRESS is present on LAN
  for ((timeout=0; timeout<=${c_timedout}; timeout++)); do
    echo -ne "Remote reboot (ping): $timeout\033[0K / ${c_timedout}s\r"
    swuk_ssh ${ip} > /dev/null 2>&1
    case ${?} in
      0)  # All done (early). Ping & SSH established
        break
      ;;
      3)  # Ping fail (Ping already waited 1s)
        :
      ;;
      2)  # Ping OK but no SSH just yet
        break
      ;;
      1)  # Ping & SSH established but issues with host file
        echo "Failed to access known hosts file ${hosts}!" >&2
        exit 1
      ;;
      *)
        :
      ;;
    esac
  done
  tput cnorm
  echo

  [[ ${timeout} -gt ${c_timedout} ]] &&
    echo "Timed out, check connection!" >&2 &&
    exit 1

  tput civis
  # Get SSH public key from the IP-ADDRESS
  for ((timeout=0; timeout<=${c_timedout}; timeout++)); do
    echo -ne "Remote reboot (ssh): $timeout\033[0K / ${c_timedout}s\r"
    swuk_ssh ${ip} > /dev/null 2>&1
    case ${?} in
      0)  # All done. Ping & SSH established
        break
      ;;
      3)  # Ping fail (Ping already waited 1s)
        :
      ;;
      2)  # Ping OK but no SSH just yet
        sleep 1
      ;;
      1)  # Ping & SSH established but issues with host file
        echo "Failed to access known hosts file ${hosts}!" >&2
        exit 1
      ;;
      *)
        :
      ;;
    esac
  done
  tput cnorm
  echo

  [[ ${timeout} -gt ${c_timedout} ]] &&
    echo "Timed out, check connection!" >&2 &&
    exit 1

  # Success message
  echo -e "\nUploaded ${c_bfiles[0]}, ${c_bfiles[1]} & ${c_bfiles[2]} to ${c_remote} @ ${ip}."

  # Display remote details
  sshpass -p ${password} scp ${username}@${ip}:/srv/www/project.txt /tmp > /dev/null 2>&1
  [ ${?} -ne 0 ] && echo "Failed to obtain remote system credentials (/srv/www/project.txt)!" >&2 && exit 1
  echo -e "\nRemote system credentials :-"
  cat /tmp/project.txt
}


################################################################################
# Opening gambit.
# Arguments ...... ${@}
# Return ......... None
# Exit code ...... Status (0=success, 1=failure)
# Shared (In) .... None
# Shared (Out) ... None
main ${@}
exit 0