#!/bin/bash

# Default values
dns_server="1.1.1.1"
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

# Check if DNS server is provided
if [ -z "$dns_server" ]; then
  echo "Error: DNS server not provided"
  echo "Usage: $0 -s <dns_server> -d <query_domain> -w <warning_time> -c <critical_time>"
  exit 1
fi

# Perform DNS resolution check and measure the time
start_time=$(date +%s%N)
result=$(dig @$dns_server $query_domain +short)
end_time=$(date +%s%N)
response_time=$(((end_time - start_time) / 1000000))
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 -q "no servers could be reached"; then
  echo "CRITICAL - DNS resolution failed for $query_domain on $dns_server | $perfdata"
  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