#!/usr/bin/env bash

#
# This script requires an iPerf3 server.
#
# How to set up the server:
# 1. `sudo apt install iperf3`
# 2. `mkdir -p /etc/iperf3`
# 3. Set a username and password (password does not matter, we're going to use key-based authentication): `S_USER=your_username S_PASSWD=your_password`
# 4. Hash your password: `echo -n "{$S_USER}$S_PASSWD" | sha256sum | awk '{ print $1 }'`
# 5. Put the output in /etc/iperf3/users.csv like this:
# mario,bf7a49a846d44b454a5d11e7acfaf13d138bbe0b7483aa3e050879700572709b
# 6. Generate your server key: `openssl genrsa -des3 -out /etc/iperf3/iperf3_server.private.pem 2048`
# 7. Generate your public key: `openssl rsa -in /etc/iperf3/iperf3_server.private.pem -outform PEM -pubout -out /etc/iperf3/iperf3.public.pem`
# 8. Generate the key for iPerf3: `openssl rsa -in /etc/iperf3/iperf3_server.private.pem -out /etc/iperf3/iperf3_client.public.pem -outform PEM`
# 9: Create the systemd service:
# cat > /etc/systemd/system/iperf3.service <<EOF
# [Unit]
# Description=iperf3 server
# After=syslog.target network.target auditd.service
#
# [Service]
# ExecStart=/usr/bin/iperf3 -s --rsa-private-key-path /etc/iperf3/server.private.pem --authorized-users-path /etc/iperf3/users.csv
#
# [Install]
# WantedBy=multi-user.target
# EOF
#
# 10. systemctl daemon-reload
# 11. systemctl enable --now iperf3
# 12. systemctl status iperf3
# 13. Run the client
#

# Warnng and critical levels are based on your specific network speed.

# Default values
SERVER=""
WARNING_LEVEL=""
CRITICAL_LEVEL=""
RSA_PUBLIC_KEY=""
USERNAME=""
PASSWORD=""
RETRY=3

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

  case $key in
  --server)
    SERVER="$2"
    shift
    shift
    ;;
  --warning)
    WARNING_LEVEL="$2"
    shift
    shift
    ;;
  --critical)
    CRITICAL_LEVEL="$2"
    shift
    shift
    ;;
  --rsa-public-key)
    RSA_PUBLIC_KEY="$2"
    shift
    shift
    ;;
  --username)
    USERNAME="$2"
    shift
    shift
    ;;
  --password)
    PASSWORD="$2"
    shift
    shift
    ;;
  --retry)
    RETRY="$2"
    shift
    shift
    ;;
  *)
    shift
    ;;
  esac
done

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

# Check if required arguments are provided
if [[ -z "$SERVER" ]] || [[ -z "$WARNING_LEVEL" ]] || [[ -z "$CRITICAL_LEVEL" ]]; then
  echo "Usage: $0 --server SERVER --warning WARNING_LEVEL --critical CRITICAL_LEVEL [--rsa-public-key RSA_PUBLIC_KEY] [--username USERNAME] [--password PASSWORD]"
  exit 3
fi

# Set IPERF3_PASSWORD environment variable
export IPERF3_PASSWORD="$PASSWORD"

# Run iperf3 command with optional arguments
for ((i = 1; i <= RETRY; i++)); do
  if [[ -n "$RSA_PUBLIC_KEY" ]] && [[ -n "$USERNAME" ]]; then
    OUTPUT=$(iperf3 -c "$SERVER" -i 1 -t 30 -f m --rsa-public-key-path "$RSA_PUBLIC_KEY" --username "$USERNAME" 2>&1)
  else
    OUTPUT=$(iperf3 -c "$SERVER" -i 1 -t 10 -f m 2>&1)
  fi

  # Check if iperf3 command failed
  if [[ $? -ne 0 ]]; then
    if [[ $OUTPUT == *"the server is busy running a test"* ]]; then
      if [[ $i -lt $RETRY ]]; then
        sleep 60
        continue
      fi
    fi
    echo -e "UNKNOWN - iperf3 command failed: $OUTPUT\n"
    exit 3
  else
    break
  fi
done

# Extract the receiver bitrate
RECEIVER_BITRATE=$(echo "$OUTPUT" | grep -Eo '[0-9]+(\.[0-9]+)? Mbits/sec' | tail -1 | awk '{print $1}')

# Prepare performance data
PERFDATA="receiver_bitrate=${RECEIVER_BITRATE}mb;${WARNING_LEVEL};${CRITICAL_LEVEL};0;"

# Check the receiver bitrate against warning and critical levels
if (($(echo "$RECEIVER_BITRATE < $CRITICAL_LEVEL" | bc -l))); then
  echo "CRITICAL - Receiver Bitrate: $RECEIVER_BITRATE Mbits/sec | $PERFDATA"
  exit 2
elif (($(echo "$RECEIVER_BITRATE < $WARNING_LEVEL" | bc -l))); then
  echo "WARNING - Receiver Bitrate: $RECEIVER_BITRATE Mbits/sec | $PERFDATA"
  exit 1
else
  echo "OK - Receiver Bitrate: $RECEIVER_BITRATE Mbits/sec | $PERFDATA"
  exit 0
fi