custom_ssh

#!/bin/bash
SSH_timeout=3
PRIMARY_HOSTNAME="Harbormaster"
SECONDARY_HOSTNAME="Monkeebutt"
JUMPBOX="BlackPearl"
command=""

sshukh () {
  output=$(\ssh "$@" 2>&1 | tee /dev/tty)
  error=$(echo $output | tail -1)
  if [[ "$error" == "Host key verification failed."* ]]; then
    host=$(cut -d'@' -f2 <<< $1)
    while true; do
      read -p "Update known_hosts? [y/n] " yn
      case $yn in
        [Yy]* ) ssh-keygen -R $host && \ssh "$@"; break;;
        [Nn]* ) break;;
        * ) echo "Please answer y or n.";;
      esac
    done
  fi
}

# Function to interpret color codes for any command output
colorize_output() {
  awk '{ gsub(/\\e/, "\x1b"); print }'
}

IsHost() {
    local hostname=$(echo "$1" | awk '{print tolower($0)}')
    local hosts_file="/etc/hosts"
    #echo "Checking if host $hostname is in hosts file."  # Debug
    if grep -wqi "\b$hostname\b" "$hosts_file"; then
        return 0
    else
        return 1
    fi
}

# Function to determine if an IP address is within a private range
IsPrivateIP() {
    local ip=$1
    if [[ $ip =~ ^10\. || $ip =~ ^172\.(1[6-9]|2[0-9]|3[0-1])\. || $ip =~ ^192\.168\. ]]; then
        return 0 # True, IP is private
    else
        return 1 # False, IP is not private
    fi
}

# Updated resolve_hostname function to use getent instead of nslookup
resolve_hostname() {
    local input="$1"
    
    # If the input is already an IP address, return it
    if [[ "$input" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
        echo "$input"
        return 0  # Return success because it's already an IP address
    fi

    # Use getent to resolve the hostname to an IP address
    local resolved_ip
    resolved_ip=$(getent hosts "$input" | awk '{ print $1 }')

    if [[ -z "$resolved_ip" ]]; then
        #echo "Failed to resolve hostname '$input'." >&2
        return 1  # Return failure
    else
        echo "$resolved_ip"
        return 0  # Return success
    fi
}

# Updated IsDifferentSubnet function with improved logic
IsDifferentSubnet() {
    local local_ip=$(hostname -I | awk '{print $1}')
    local target_ip=$(resolve_hostname "$1")

    if [[ -z "$target_ip" ]]; then
        return 0  # Assume different subnet if no IP could be found
    fi

    # Extract the first three octets (subnet) for both local and target IPs
    local local_subnet=${local_ip%.*}
    local target_subnet=${target_ip%.*}

    # Compare the subnets
    if [[ "$local_subnet" == "$target_subnet" ]]; then
        return 1  # False, same subnet
    else
        return 0  # True, different subnet
    fi
}


# Function to extract the username from the destination string
extract_username() {
    if [[ "$1" =~ @ ]]; then
        echo "$1" | awk -F'@' '{print $1}'
    else
        echo ""
    fi
}

if [ $# -eq 0 ]; then
    case $HOSTNAME in
        $PRIMARY_HOSTNAME)
            #echo "Defaulting to $destination"
            destination="$SECONDARY_HOSTNAME"
            ;;
        $SECONDARY_HOSTNAME)
            #echo "Defaulting to $destination"
            destination="$PRIMARY_HOSTNAME"
            ;;
        *)
            echo "No destination specified and unknown hostname." >&2
            exit 1
            ;;
    esac
elif [ $# -gt 0 ]; then
    input_host="$1"
    if [[ "$input_host" =~ @ ]]; then
        input_host=$(echo "$input_host" | awk -F'@' '{print $2}')
    fi

    # Resolve the hostname, but don't exit if it fails yet
    destination=$(resolve_hostname "$input_host")
    if [ $? -ne 0 ]; then
    # Check if the input is a recognized host
    if IsHost "$input_host"; then
        destination="$input_host"
        echo "Using local hostname: $destination"
        username=$(extract_username "$1")
        user_prefix=""
        if [ -n "$username" ]; then
            user_prefix="$username@"
        fi
        if [ $# -gt 1 ]; then
            shift  # Remove the first argument which is the destination
            command="$@"
        else
            command=""
        fi
    else
        # If not a recognized host, treat input as a command
        # Set destination based on the current host
        case $(hostname -s) in
          $PRIMARY_HOSTNAME)
            destination="$SECONDARY_HOSTNAME"
            ;;
          $SECONDARY_HOSTNAME)
            destination="$PRIMARY_HOSTNAME"
            ;;            *)
                echo "Unknown current host. Exiting." >&2
                exit 1
                ;;
        esac

        # Set the command to the input value
        command="$input_host"
        #echo "Defaulting to destination: $destination with command: $command" >&2
    fi
    else
        # If hostname is resolved, prepare to use it as the destination
        username=$(extract_username "$1")
        user_prefix=""
        if [ -n "$username" ]; then
            user_prefix="$username@"
        fi
        if [ $# -gt 1 ]; then
            shift  # Remove the first argument which is the destination
            command="$@"
        else
            command=""
        fi
    fi
else
    echo "No arguments specified." >&2
    exit 1
fi

if IsDifferentSubnet "$destination"; then
    echo "Using jump box: ssh -J $JUMPBOX ${destination} ${command}" >&2
    if [[ -n $command ]]; then
        sshukh  -t -o ConnectTimeout=$SSH_timeout -J "$JUMPBOX" "${user_prefix}${destination}" "${command}" | colorize_output
    else
        sshukh -t -o ConnectTimeout=$SSH_timeout -J "$JUMPBOX" "${user_prefix}${destination}"
    fi

else
    if [[ -n $command ]]; then
        sshukh -t -o ConnectTimeout=$SSH_timeout "${user_prefix}${destination}" "${command}" | colorize_output
    else
        sshukh -t -o ConnectTimeout=$SSH_timeout "${user_prefix}${destination}"
    fi
fi