#!/bin/bash

# Default values
HOST=""
PORT=22
PASSWORD=""
SSH_KEY=""
USERNAME=""
TIMEOUT=10
WARNING=3
CRITICAL=5
MODE="auth"

# Usage instructions
usage() {
  echo "Usage: $0 -H <host> -u <username> [-p <password> | -k <ssh_key>] [-P <port>] [-t <timeout>] [-w <warning>] [-c <critical>] [-m <mode>]"
  echo
  echo "  -H <host>       : SSH server hostname or IP address"
  echo "  -u <username>   : SSH username"
  echo "  -p <password>   : SSH password (optional if using SSH key)"
  echo "  -k <ssh_key>    : Path to SSH private key (optional. \$HOME/.ssh/id_rsa is used by default)"
  echo "  -P <port>       : SSH port (default: 22)"
  echo "  -t <timeout>    : Connection timeout in seconds (default: 10)"
  echo "  -w <warning>    : Warning threshold for connection time in seconds (default: 3)"
  echo "  -c <critical>   : Critical threshold for connection time in seconds (default: 5)"
  echo "  -m <mode>       : Mode for testing ('auth' or 'conn', default: 'auth'. If one mode fails, try the other)"
  echo
  exit 3
}

# Parse arguments
while getopts "H:P:p:k:u:t:w:c:m:" opt; do
  case $opt in
    H) HOST="$OPTARG" ;;
    P) PORT="$OPTARG" ;;
    p) PASSWORD="$OPTARG" ;;
    k) SSH_KEY="$OPTARG" ;;
    u) USERNAME="$OPTARG" ;;
    t) TIMEOUT="$OPTARG" ;;
    w) WARNING="$OPTARG" ;;
    c) CRITICAL="$OPTARG" ;;
    m) MODE="$OPTARG" ;;
    *) usage ;;
  esac
done

# Check if required arguments are provided
if [ -z "$HOST" ] || [ -z "$USERNAME" ]; then
  usage
fi

# Use default SSH key if not specified
if [ -z "$SSH_KEY" ]; then
  SSH_KEY="$HOME/.ssh/id_rsa"
fi

# Check connection
START_TIME=$(date +%s%N)
if [ -z "$PASSWORD" ]; then
  if [ "$MODE" == "auth" ]; then
    OUTPUT=$(ssh -q -o ConnectTimeout=$TIMEOUT -o StrictHostKeyChecking=no -o BatchMode=yes -i $SSH_KEY -p $PORT $USERNAME@$HOST "echo success" 2>&1); rc=$?
  else
    OUTPUT=$(ssh -q -o ConnectTimeout=$TIMEOUT -o StrictHostKeyChecking=no -o BatchMode=yes -o ProxyCommand="nc -w $TIMEOUT %h %p" -i $SSH_KEY -p $PORT $USERNAME@$HOST 2>&1); rc=$?
  fi
else
  if [ "$MODE" == "auth" ]; then
    OUTPUT=$(sshpass -p "$PASSWORD" ssh -q -o ConnectTimeout=$TIMEOUT -o StrictHostKeyChecking=no -o BatchMode=yes -p $PORT $USERNAME@$HOST "echo success" 2>&1); rc=$?
  else
    OUTPUT=$(sshpass -p "$PASSWORD" ssh -q -o ConnectTimeout=$TIMEOUT -o StrictHostKeyChecking=no -o BatchMode=yes -o ProxyCommand="nc -w $TIMEOUT %h %p" -p $PORT $USERNAME@$HOST 2>&1); rc=$?
  fi
fi
END_TIME=$(date +%s%N)
ELAPSED_TIME=$(echo "scale=3; ($END_TIME - $START_TIME) / 1000000000" | bc)
ELAPSED_TIME_MS=$(echo "scale=0; ($END_TIME - $START_TIME) / 1000000" | bc)

# Check result
perfdata="response_time=${ELAPSED_TIME_MS}ms;${WARNING};${CRITICAL};0;;"
if [[ $MODE == "auth" && $OUTPUT == "success" ]] || [[ $MODE == "conn" && $rc -eq 0 ]]; then
  if (( $(echo "$ELAPSED_TIME > $CRITICAL" | bc -l) )); then
    echo "CRITICAL - connected to $HOST in $ELAPSED_TIME seconds | $perfdata"
    exit 2
  elif (( $(echo "$ELAPSED_TIME > $WARNING" | bc -l) )); then
    echo "WARNING - connected to $HOST in $ELAPSED_TIME seconds | $perfdata"
    exit 1
  else
    echo "OK - connected to $HOST in $ELAPSED_TIME seconds | $perfdata"
    exit 0
  fi
else
  echo -e "CRITICAL - failed to connect to $HOST. Return code was $rc\n$OUTPUT"
  exit 2
fi