#!/usr/bin/env bash

SERVER_ADDRESS=""
SECRET=""
REALM=""
PORT=3478
PROTOCOL="udp"
TIMEOUT=10
LOST_PACKETS_WARN=3
LOST_PACKETS_CRIT=5
SEND_DROPPED_WARN=3
SEND_DROPPED_CRIT=5
AVG_RTT_WARN=100
AVG_RTT_CRIT=200
REQUESTED_RELAY_ADDR="false"


while getopts "s:S:u:r:p:w:c:W:C:R:T:x:" opt; do
  case $opt in
  s)
    SERVER_ADDRESS="$OPTARG"
    ;;
  S)
    SECRET="$OPTARG"
    ;;
  r)
    REALM="$OPTARG"
    ;;
  p)
    PORT="$OPTARG"
    ;;
  w)
    LOST_PACKETS_WARN="$OPTARG"
    ;;
  c)
    LOST_PACKETS_CRIT="$OPTARG"
    ;;
  W)
    SEND_DROPPED_WARN="$OPTARG"
    ;;
  C)
    SEND_DROPPED_CRIT="$OPTARG"
    ;;
  R)
    AVG_RTT_WARN="$OPTARG"
    ;;
  T)
    AVG_RTT_CRIT="$OPTARG"
    ;;
  x)
    REQUESTED_RELAY_ADDR="-x $OPTARG"
    ;;
  *)
    echo "Usage: $0 -s SERVER_ADDRESS -S SECRET -r REALM [-p PORT] [-P PROTOCOL] [-t TIMEOUT] [-w LOST_PACKETS_WARN] [-c LOST_PACKETS_CRIT] [-W SEND_DROPPED_WARN] [-C SEND_DROPPED_CRIT] [-R AVG_RTT_WARN] [-T AVG_RTT_CRIT] [-x REQUESTED_RELAY_ADDR]"
    exit -1
    ;;
  esac
done

# Check if required options are set
if [[ -z "$SERVER_ADDRESS" || -z "$SECRET" || -z "$REALM" ]]; then
  echo "Usage: $0 -s SERVER_ADDRESS -S SECRET -r REALM [-p PORT] [-P PROTOCOL] [-t TIMEOUT] [-w LOST_PACKETS_WARN] [-c LOST_PACKETS_CRIT] [-W SEND_DROPPED_WARN] [-C SEND_DROPPED_CRIT] [-R AVG_RTT_WARN] [-T AVG_RTT_CRIT] [-x REQUESTED_RELAY_ADDR]"
  exit -1
fi

if ! command -v turnutils_uclient &>/dev/null; then
  echo "UNKNOWN - turnutils_uclient not found! Please install coturn. Make sure to run \`sudo apt install coturn; sudo service coturn stop; sudo systemctl disable coturn\`"
  exit -1
fi

if ! command -v stun &>/dev/null; then
  echo "UNKNWON - stun not found! Please install stun-client"
  exit -1
fi

# Fetch the user's public IP using the coturn server as a STUN server
PEER_ADDRESS_OUT=$(stun "$SERVER_ADDRESS" -p "$PORT" -v 1 2>&1)
PEER_ADDRESS=$(echo "$PEER_ADDRESS_OUT" | grep "MappedAddress" | awk -F'[ =:]+' '{print $2}')

if [[ -z "$PEER_ADDRESS" ]]; then
  echo "UNKNOWN - failed to fetch this host's public IP using the coturn server as a STUN server."
  echo "$PEER_ADDRESS_OUT"
  exit -1
fi

TURNUTILS_OUTPUT=$(turnutils_uclient "$SERVER_ADDRESS" -s -W "$SECRET" -r "$REALM" -p "$PORT" -e "$PEER_ADDRESS" $REQUESTED_RELAY_ADDR -B -y 2>&1)

if [ $? -eq 0 ]; then
  LOST_PACKETS=$(echo "$TURNUTILS_OUTPUT" | grep -m1 "Total lost packets" | awk -F'[(%)]' '{print $2"%"}' | tr -d '%' | cut -d. -f1)
  SEND_DROPPED=$(echo "$TURNUTILS_OUTPUT" | grep -m1 "Total lost packets" | awk -F'[(%)]' '{print $5"%"}' | tr -d '%' | cut -d. -f1)
  AVG_RTT=$(echo "$TURNUTILS_OUTPUT" | grep -m1 "Average round trip delay" | awk '{print $7}' | cut -d. -f1)
  AVG_JITTER=$(echo "$TURNUTILS_OUTPUT" | grep -m1 "Average jitter" | awk '{print $5}' | cut -d. -f1)

  EXIT_CODE=0
  STATUS="OK"

  if [ "$LOST_PACKETS" -ge "$LOST_PACKETS_CRIT" ] || [ "$SEND_DROPPED" -ge "$SEND_DROPPED_CRIT" ] || [ "$AVG_RTT" -ge "$AVG_RTT_CRIT" ]; then
    STATUS="CRITICAL"
    EXIT_CODE=2
  elif [ "$LOST_PACKETS" -ge "$LOST_PACKETS_WARN" ] || [ "$SEND_DROPPED" -ge "$SEND_DROPPED_WARN" ] || [ "$AVG_RTT" -ge "$AVG_RTT_WARN" ]; then
    STATUS="WARNING"
    EXIT_CODE=1
  fi

  echo "$STATUS - connected to TURN server $SERVER_ADDRESS:$PORT. lost_packets=$LOST_PACKETS% send_dropped=$SEND_DROPPED% avg_rtt=${AVG_RTT}ms avg_jitter=${AVG_JITTER}ms | lost_packets_percent=$LOST_PACKETS;$LOST_PACKETS_WARN;$LOST_PACKETS_CRIT;0;100 send_dropped_percent=$SEND_DROPPED;$SEND_DROPPED_WARN;$SEND_DROPPED_CRIT;0;100 avg_rtt=${AVG_RTT}ms;$AVG_RTT_WARN;$AVG_RTT_CRIT;0; avg_jitter=${AVG_JITTER}ms;0;" exit $EXIT_CODE
else
  echo "CRITICAL - failed to connect to TURN server:"
  echo "$TURNUTILS_OUTPUT"
  exit 2
fi