Advanced SSH

# SourceMe
# Recomended to set up passwordless ssh authentication https://linuxize.com/post/how-to-setup-passwordless-ssh-login/
# Copy the .pub key to both login.ansatt.ntnu.no and physics cluster workstation

function sshtun() {
  # Replace $USER with your NTNU username
  local USERNAME=$USER
  if [[ -z "$USERNAME" ]]; then
    echo "USERNAME is not defined"
    return 1
  fi
  local tunnel_sessions="${HOME}/.ssh/ssh-tunnels-created"
  if [[ ${1} == "help" ]] || [[ -z "$1" ]]; then
    local __help="
    Usage: ${FUNCNAME[0]} hostname
    where hostname must be a FQDN hostname! (eg. igor.phys.ntnu.no)
    This function will create a multiplexed ssh tunnel through one of the 3 hosts at NTNU
    after which will establish an ssh with X11 forwarding connection to your host using 
    the tunnel. This way the connections, especially those requiring GUIs, will be really
    fast and responsive.
    Each time a new tunnel will be created even if you're trying to connect to the same 
    host. To connect multiple times to the same host, follow the instructions on the screen
    after establishing a connection or, check the file \"$tunnel_sessions\"
    "
    echo "$__help"
    return 0
  fi
  host ${1}
  if [ $? -eq 0 ]; then
    local options=( "login.ansatt.ntnu.no" "login.phys.ntnu.no" "login.stud.ntnu.no" )
    # Print the menu options
    echo "Select the tunneling host"
    for ((i=0; i<${#options[@]}; i++)); do
      echo "$((i+1)). ${options[$i]}"
    done

    # Prompt the user to enter a choice
    read -p "Enter your choice (1-${#options[@]}) (1): " choice
    # Default option [1]; Change to another one if you want
    if [[ -z "$choice" ]]; then
      choice=1
    fi
    echo "Choice was: $choice"
    # Validate the input
    re='^[0-9]+$'
    if ! [[ $choice =~ $re ]] || (( choice < 1 || choice > ${#options[@]} )); then
      return 1
    fi
      # Process the chosen option
    case $choice in
      1)
        t_host=tun-ansatt
        m_host=login.ansatt.ntnu.no
        ;;
      2)
        t_host=tun-phys
        m_host=login.phys.ntnu.no
        ;;
      3)
        t_host=tun-stud
        m_host=login.stud.ntnu.no
        ;;
      *)
        echo "Not a valid choice. Exiting"
        return 1
        ;;
    esac
    
    echo "Connecting to ${1} through ${m_host} with user ${USERNAME}"
    local PORT=$(( ((RANDOM<<15)|RANDOM) % 63001 + 2000 ))
    local counter=0
    while [[ $(netstat -ant -p TCP | grep ${PORT}) -ne 0 ]]; do
      echo "Trying new port ..."
      local PORT=$(( ((RANDOM<<15)|RANDOM) % 63001 + 2000 ))
      ((counter++))
      if [[ $counter -gt 5 ]]; then
        echo "Couldn't find a free port"
        return 1
      fi
      sleep 1
    done

      local SSHCPath="${HOME}/.ssh/${USERNAME}@${m_host}:${PORT}-${1}"
      ssh -4 -f -N -T -M \
          -o ExitOnForwardFailure=yes \
          -o ControlMaster=autoask \
          -o ControlPersist=10m \
          -o ControlPath="${SSHCPath}" \
          -L 127.0.0.1:${PORT}:${1}:22 ${USERNAME}@${m_host}
      
      local exit_code=$?
      if [[ $exit_code -eq 0 ]]; then 
        sleep 3
        
        echo -e "\n****************************************************" >> $tunnel_sessions
        echo "* @ $(date +%F/%T)" >> $tunnel_sessions
        echo "****************************************************" >> $tunnel_sessions
        echo -e "You can reconnect/reuse the tunnel to ${1} with:\n ssh -4 -X  ${USERNAME}@127.0.0.1 -p ${PORT}" | tee -a "$tunnel_sessions"
        echo "Check ${tunnel_sessions} for tunnels opened"
        echo "FORWARDING added to port: ${PORT}"
        echo -e "Check for ControlMaster connections status:\n ssh -O check -S \"${SSHCPath}\" ${1}" | tee -a "$tunnel_sessions"
        echo -e "Stop the tunnels from accepting connections:\n ssh -O stop -S \"${SSHCPath}\" ${1}" | tee -a "$tunnel_sessions"
        ssh  -4 -X ${USERNAME}@127.0.0.1 -p${PORT}
      else
        echo "ERROR: Couldn't create multiplexed connection!" 
        retun 1
      fi
  else
    echo -e "Hostname not found. Must be a FQDN hostname! (eg. igor.phys.ntnu.no)"
    echo "Type \"${FUNCNAME[0]} help\" for more info."
  fi
}