#!/bin/bash

# Default values
dns_server=""
query_domain=""
warning_time=200
critical_time=250

# Parse flag arguments
while getopts "s:d:w:c:" flag; do
  case "${flag}" in
  s) dns_server="${OPTARG}" ;;
  d) query_domain="${OPTARG}" ;;
  w) warning_time="${OPTARG}" ;;
  c) critical_time="${OPTARG}" ;;
  *)
    echo "Usage: $0 -s <dns_server> -d <query_domain> -w <warning_time> -c <critical_time>"
    exit 1
    ;;
  esac
done

if [[ -z $query_domain && ${query_domain+x} ]]; then
  echo "Error: query domain not provided."
  echo "Usage: $0 -s <dns_server> -d <query_domain> -w <warning_time> -c <critical_time>"
  exit 3
fi

# Perform DNS resolution check and measure the time
start_time=$(date +%s)

if [[ -z "$dns_server" ]]; then
  # The DNS server was not provided.
  result=$(dig "$query_domain" +short)
  dns_server="the default server"
else
  result=$(dig @$dns_server "$query_domain" +short)
fi

end_time=$(date +%s)
response_time=$((end_time - start_time))
perfdata="response_time=${response_time}ms;${warning_time};${critical_time};0;"

# Check response time against warning and critical levels
if [ -z "$result" ] || echo "$result" | grep -Eq "(no servers could be reached|communications error)"; then
  echo "CRITICAL - DNS resolution failed for $query_domain on $dns_server"
  exit 2
elif [ "$critical_time" -gt 0 ] && [ "$response_time" -gt "$critical_time" ]; then
  echo "CRITICAL - DNS resolution took ${response_time}ms for $query_domain on $dns_server | $perfdata"
  exit 2
elif [ "$warning_time" -gt 0 ] && [ "$response_time" -gt "$warning_time" ]; then
  echo "WARNING - DNS resolution took ${response_time}ms for $query_domain on $dns_server | $perfdata"
  exit 1
else
  echo "OK - resolved $query_domain to $result using $dns_server in ${response_time}ms | $perfdata"
  exit 0
fi