#!/bin/bash

#
# File .......... swuk_ssh
# Author ........ Steve Haywood
# Website ....... http://www.spacewire.co.uk
# Project ....... Common (SpaceWire UK Tutorial)
# Date .......... 18 Jun 2025
# Version ....... 1.0
# Description ...
#   Enable pathway to PetaLinux running on the Zedboard.
#

# Strict
#set -euo pipefail


################################################################################
# Main function.
# Arguments ...... ${@}
# Return ......... None
# Exit code ...... Status (0=success, 1=failure)
# Shared (In) .... None
# Shared (Out) ... None
main()
{
  # Declare 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
                                # Str: Known hosts file
  local -r  hosts="${HOME}/.ssh/known_hosts"
                                # Str: IP validation regular expression
  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]))$'

  # 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)

  # Display help information
  if [[ " ${argv[*]} " =~ " ${c_opthelp} " ]]; then
    echo "Usage: $(basename ${0}) IP-ADDRESS... [OPTION]..."
    echo "Establish a secure shell connection 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 IP-ADDRESS is present on LAN
  ping -W 1 -c 1 ${ip} > /dev/null 2>&1
  [ ${?} -ne 0 ] &&
    echo "No ping response ${ip}, check connection!" >&2 &&
    exit 3

  # Get SSH public key from the IP-ADDRESS
  key=$(ssh-keyscan -H ${ip} 2> /dev/null)
  if [[ -n ${key} ]]; then
    # Remove all previous keys belonging to the IP ADDRESS
    ssh-keygen -R ${ip} -f ${hosts} > /dev/null 2>&1
    [ ${?} -ne 0 ] &&
      echo "Failed to remove ${ip} from Known Hosts file ${hosts}!" >&2 &&
      exit 1
    # Add new key belonging to the IP ADDRESS
    echo ${key} >> ${hosts}
    [ ${?} -ne 0 ] &&
      echo "Failed to add ${ip} to Known Hosts file ${hosts}!" >&2 &&
      exit 1
  else
    echo "Failed to get SSH public key from ${ip}!" >&2
    exit 2
  fi

  # Success message
  echo -e "\nEstablished SSH connection with ${ip}."
}


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