#!/usr/bin/env bash

function print_help {
  echo "Usage: $0 --server SERVER --warning WARNING --critical CRITICAL --rsa-key RSA_KEY --username USERNAME --password PASSWORD"
  echo ""
  echo "Arguments:"
  echo "  --server   The iperf3 server to test against"
  echo "  --warning  Warning level for the speed test in Mbits/s"
  echo "  --critical Critical level for the speed test in Mbits/s"
  echo "  --rsa-key  Path to the RSA public key for authentication"
  echo "  --username Username for authentication"
  echo "  --password Password for authentication"
  echo ""
  exit 3
}

if ! command -v iperf3 &>/dev/null; then
  echo "UNKNOWN - iperf3 is not installed"
  exit 3
fi

if ! command -v jq &>/dev/null; then
  echo "UNKNOWN - jq is not installed"
  exit 3
fi

if [[ $# -eq 0 ]]; then
  print_help
fi

while [[ $# -gt 0 ]]; do
  key="$1"

  case $key in
  --server)
    SERVER="$2"
    shift
    shift
    ;;
  --warning)
    WARNING="$2"
    shift
    shift
    ;;
  --critical)
    CRITICAL="$2"
    shift
    shift
    ;;
  --rsa-key)
    RSA_KEY="$2"
    shift
    shift
    ;;
  --username)
    USERNAME="$2"
    shift
    shift
    ;;
  --password)
    PASSWORD="$2"
    shift
    shift
    ;;
  *)
    print_help
    ;;
  esac
done

if [[ -z "$SERVER" ]] || [[ -z "$WARNING" ]] || [[ -z "$CRITICAL" ]] || [[ -z "$RSA_KEY" ]] || [[ -z "$USERNAME" ]] || [[ -z "$PASSWORD" ]]; then
  echo "UNKNOWN - All arguments are required"
  print_help
fi

export IPERF3_PASSWORD="$PASSWORD"

RESULT=$(iperf3 -c $SERVER -u -b 0 -J -t 10 --rsa-public-key-path $RSA_KEY --username $USERNAME 2>&1)
if [[ $? -ne 0 ]]; then
  echo "UNKNOWN - iperf3 command failed: $RESULT"
  exit 3
fi

if ! echo "$RESULT" | jq -e '.end.sum.bits_per_second' >/dev/null; then
  echo "UNKNOWN - Invalid JSON output from iperf3: $RESULT"
  exit 3
fi

STREAM_RESULT=$(echo "$RESULT" | jq '.end.streams[] | select(.udp.sender == (!0)) | .udp.bits_per_second')
SPEED_RAW=$(echo "$STREAM_RESULT" | jq 'if length > 0 then . / 1000000 else 0 end')
SPEED=$(printf "%.0f" "$SPEED_RAW")

if (($(echo "$SPEED < $CRITICAL" | bc -l))); then
  echo "CRITICAL - Speed: ${SPEED} Mbits/s"
  exit 2
elif (($(echo "$SPEED < $WARNING" | bc -l))); then
  echo "WARNING - Speed: ${SPEED} Mbits/s"
  exit 1
else
  echo "OK - Speed: ${SPEED} Mbits/s"
  exit 0
fi