#!/bin/bash

function display_help {
  echo "Usage: $0 [options]"
  echo
  echo "Options:"
  echo "-i, Wireguard interface name. Identical to what you would put in wg-quick."
  echo "-s, How long to wait for a handshake in seconds, after which the script will fail. (default: $HANDSHAKE_SEC_DEFAULT)"
  echo "-w, Handshake time warning level in seconds. (default: $WARNING_LEVEL_DEFAULT)"
  echo "-c, Handshake time critical level in seconds. (default: $CRITICAL_LEVEL_DEFAULT)"
  echo "-h, Display this help menu"
  echo
}

while getopts ":i:s:w:c:h" opt; do
  case ${opt} in
  i)
    WG_INTERFACE=$OPTARG
    ;;
  s)
    HANDSHAKE_SEC=$OPTARG
    ;;
  w)
    WARNING_LEVEL=$OPTARG
    ;;
  c)
    CRITICAL_LEVEL=$OPTARG
    ;;
  h)
    display_help
    exit 0
    ;;
  \?)
    echo "Invalid option: -$OPTARG" 1>&2
    display_help
    exit 1
    ;;
  :)
    echo "Option -$OPTARG requires an argument." 1>&2
    display_help
    exit 1
    ;;
  esac
done

# Set default values
HANDSHAKE_SEC=${HANDSHAKE_SEC:-20}
WARNING_LEVEL=${WARNING_LEVEL:-5}
CRITICAL_LEVEL=${CRITICAL_LEVEL:-10}

# Check if WG_INTERFACE is set
if [ -z "$WG_INTERFACE" ]; then
  echo "Error: -i (interface) option is required" 1>&2
  display_help
  exit 1
fi

if ! command -v wg-quick &>/dev/null; then
  echo "UNKNOWN - wg-quick could not be found"
  exit 3
fi

function cleanup {
  wg-quick down "$WG_INTERFACE" >/dev/null 2>&1
}
trap cleanup EXIT

# Try to establish a connection
WG_UP_OUTPUT=$(wg-quick up "$WG_INTERFACE" 2>&1)

# Check if the connection was successful
WG_STATUS=$(wg show)

if [ -z "$WG_STATUS" ]; then
  wg-quick down "$WG_INTERFACE" >/dev/null 2>&1 # be extra careful
  echo "CRITICAL - Unable to establish a connection to the Wireguard server. Output from wg-quick up:"
  echo "$WG_UP_OUTPUT"
  exit 2
fi

# Start the timer
START_TIME=$(date +%s%N)

# Wait for handshake
for i in $(seq 1 $(expr $HANDSHAKE_SEC \* 100)); do
  HANDSHAKE_TIME=$(wg show $WG_INTERFACE latest-handshakes | awk '{print $2}')
  CURRENT_TIME=$(date +%s)
  if [ $(expr $CURRENT_TIME - $HANDSHAKE_TIME) -le $HANDSHAKE_SEC ]; then
    break
  fi
  if [ $i -eq $(expr $HANDSHAKE_SEC \* 100) ]; then
    echo "CRITICAL - No handshakes within the $HANDSHAKE_SEC seconds"
    wg-quick down $WG_INTERFACE >/dev/null 2>&1
    exit 2
  fi
  sleep 0.01
done

END_TIME=$(date +%s%N)
ELAPSED_TIME=$(expr $END_TIME - $START_TIME)
ELAPSED_TIME=$(echo "scale=2; $ELAPSED_TIME/1000000000" | bc)
ELAPSED_TIME=$(printf "%0.2f" $ELAPSED_TIME)

# Count peers
PEER_COUNT=$(wg show "$WG_INTERFACE" peers | wc -l)

# Check if elapsed time exceeds warning or critical levels
if (($(echo "$ELAPSED_TIME > $CRITICAL_LEVEL" | bc -l))); then
  echo "CRITICAL - Connection time $ELAPSED_TIME seconds exceeds critical level of $CRITICAL_LEVEL seconds | time=${ELAPSED_TIME}s;${WARNING_LEVEL};${CRITICAL_LEVEL};0;${CRITICAL_LEVEL} peers=${PEER_COUNT};;0;"
  exit 2
elif (($(echo "$ELAPSED_TIME > $WARNING_LEVEL" | bc -l))); then
  echo "WARNING - Connection time $ELAPSED_TIME seconds exceeds warning level of $WARNING_LEVEL seconds | time=${ELAPSED_TIME}s;${WARNING_LEVEL};${CRITICAL_LEVEL};0;${CRITICAL_LEVEL} peers=${PEER_COUNT};;0;"
  exit 1
fi

# Close connection
wg-quick down "$WG_INTERFACE" >/dev/null 2>&1

# Output metrics
echo "OK - connection to the Wireguard server was established | time=${ELAPSED_TIME}s;${WARNING_LEVEL};${CRITICAL_LEVEL};0;${CRITICAL_LEVEL} peers=${PEER_COUNT};;0;"
exit 0