#!/usr/bin/env bash

# Run global DNS queries for a domain, in parallel! 🦄
#
# This script is useful for testing DNS propagation and performance.
# 
# Setup
# This script requires the 'dig' command to be installed.
# On macOS, you can install dig with 'brew install bind'.
# On Ubuntu/Debian, you can install dig with 'sudo apt-get install dnsutils'.
#
# Usage:
#   ./dns_global_check.sh <domain> [<record type>]
#
# Examples:
#   ./dns_global_check.sh danlevy.net         # defaults to A record
#   ./dns_global_check.sh danlevy.net MX      # queries MX records
#
# Author:  Dan Levy / https://danlevy.net / @justsml
#
# License: Yolo 🫶
#
# Note: Public DNS resolvers are subject to frequent changes, blocks or rate-limits on foreign IPs.

# Initialize parameters
VERBOSE=0
DOMAIN=""
TYPE=""
TIMEOUT=""
MAX_CONCURRENT=""

# Process all arguments
while [ $# -gt 0 ]; do
  case "$1" in
    --verbose|-v)
      VERBOSE=1
      ;;
    *)
      if [ -z "$DOMAIN" ]; then
        DOMAIN="$1"
      elif [ -z "$TYPE" ]; then
        TYPE="$1"
      elif [ -z "$TIMEOUT" ]; then
        TIMEOUT="$1"
      elif [ -z "$MAX_CONCURRENT" ]; then
        MAX_CONCURRENT="$1"
      fi
      ;;
  esac
  shift
done

if [ -z "$DOMAIN" ]; then
  echo "Usage: $0 [--verbose|-v] <domain> [<record type>] [<timeout>] [<max concurrent>]"
  echo "Example: $0 danlevy.net"
  echo "Example: $0 danlevy.net MX"
  echo "Example: $0 --verbose danlevy.net TXT 20 10"
  echo "Example: $0 danlevy.net A --verbose"
  echo "Example: MAX_CONCURRENT=20 $0 danlevy.net A"
  exit 1
fi

# Default record type to 'A' if none provided
if [ -z "$TYPE" ]; then
  TYPE="A"
fi

# Timeout in seconds for each DNS query
TIMEOUT="${TIMEOUT:-8}"
# Limit of concurrent lookups
MAX_CONCURRENT=${MAX_CONCURRENT:-10}

# Function to load DNS servers from file
load_dns_servers() {
  local custom_file="$HOME/.dns-resolvers.txt"
  if [ -f "$custom_file" ] && [ -r "$custom_file" ]; then
    if [ "$VERBOSE" = "1" ]; then
      printf "Loading DNS servers from %s\n" "$custom_file"
    fi
    # Read custom servers, skip empty lines and comments
    while IFS= read -r line; do
      line=$(echo "$line" | sed 's/#.*//;s/^[[:space:]]*//;s/[[:space:]]*$//')
      if [ ! -z "$line" ]; then
        SERVERS+=("$line")
      fi
    done < "$custom_file"
    return 0
  fi
  return 1
}

# Before defining default SERVERS array, try loading custom file
if ! load_dns_servers; then
  # Default list of DNS servers and labels (IP (Label))
  SERVERS=(
    # IMPORTANT:
    #   Public DNS resolvers are subject to frequent changes, blocks or rate-limits on foreign IPs.
    #   Feel free to add or remove servers as needed.
    "1.1.1.1 (Cloudflare 1)"
    "1.0.0.1 (Cloudflare 2)"
    "8.8.8.8 (Google DNS 1)"
    "8.8.4.4 (Google DNS 2)"
    "94.140.14.14 (AdGuard 1)"
    "94.140.14.15 (AdGuard 2)"
    "85.214.28.183 (Digitalcourage - Germany 1)"
    "135.125.237.69 (OVH - Germany 2)"
    "91.26.116.18 (Deutsche Telekom - Germany 3)"
    "223.5.5.148 (AliDNS - China 1)"
    "223.5.5.0 (AliDNS - China 2)"
    "223.5.5.84 (AliDNS - China 3)"
    "89.234.93.210 (Digiweb - Ireland 1)"
    "80.93.18.196 (Digiweb - Ireland 2)"
    "158.43.192.1 (Verizon - London 1)"
    "195.21.13.234 (GTT - London 2)"
    "195.186.1.111 (Swisscom - Switzerland 1)"
    "81.7.255.3 (Swisscom - Switzerland 2)"
    "170.64.147.31 (OVH - France 1)"
    "54.37.30.59 (OVH - France 2)"
    "178.255.79.70 (Telecom Italia - Italy 1)"
    # "89.96.49.60 (Fastweb - Italy 2)"
    "197.155.92.20 (Liquid Teleco - Kenya 1)"
    # "197.248.131.203 (Safaricom - Kenya 2)"
    # "197.253.8.109 (Mainone - Nigeria 1)"
    # "196.25.1.11 (Telkom SA - South Africa 1)"
    # "196.25.1.9 (Telkom SA - South Africa 2)"
    "115.99.172.26 (BSNL - India 1)"
    "115.98.145.112 (Hathway - India 2)"
    "202.136.162.11 (M1 - Singapore 1)"
    "202.136.162.12 (M1 - Singapore 2)"
    "172.104.90.123 (Akamai - Japan 1)" 
    "153.156.93.5 (NTT - Japan 2)"
    "210.163.158.224 (NTT - Japan 3)"
    "168.126.63.1 (KT - South Korea 1)"
    "168.126.63.2 (KT - South Korea 2)"
  )
fi

# Function to get current time in milliseconds
get_time() {
  if [[ "$OSTYPE" == "darwin"* ]]; then
    # On macOS, use perl since date doesn't support nanoseconds
    perl -MTime::HiRes -e 'printf("%.0f\n", Time::HiRes::time() * 1000)'
  else
    # On Linux, use date with millisecond precision
    date +%s%3N
  fi
}

current_jobs=0
total_servers=${#SERVERS[@]}
script_start_time=$(get_time)

# Initialize counters
successful_queries=0
failed_queries=0

# Function to get safe filename from domain
get_safe_filename() {
  echo "$1" | tr -c '[:alnum:]' '_'
}

# Function to perform a DNS lookup
dns_query() {
  local server_ip="$1"
  local server_label="$2"
  local domain="$3"
  local record_type="$4"
  local timeout="$5"
  local start_time
  local end_time
  local query_status
  local resolved_ips
  local first_resolved
  local safe_domain=$(get_safe_filename "$domain")
  local tmp_prefix="/tmp/dns_check_${safe_domain}_${record_type}_$$"

  start_time=$(get_time)
  resolved_ips=$(dig +tries=1 +short +time="$timeout" @"$server_ip" "$domain" "$record_type" 2>/dev/null)
  query_status=$?
  end_time=$(get_time)
  local duration_ms=$((end_time - start_time))

  if [ $query_status -eq 0 ] && [ ! -z "$resolved_ips" ]; then
    if [ "$VERBOSE" = "1" ]; then
      printf " ✅ %d ms @ %s - %s\n" "$duration_ms" "$server_label" "$server_ip"
      printf "   └─ Resolved: %s\n" "$(echo "$resolved_ips" | tr '\n' ' ')"
    fi
    echo "success" > "${tmp_prefix}_${server_ip}"
    echo "$resolved_ips" > "${tmp_prefix}_${server_ip}.resolved"
  else
    printf " ❌ %s [%s] - %d ms\n" "$server_label" "$server_ip" "$duration_ms"
    echo "fail" > "${tmp_prefix}_${server_ip}"
  fi
}

# Loop through each DNS server and run lookups in parallel
for entry in "${SERVERS[@]}"; do
  IP=$(echo "$entry" | awk '{print $1}')
  LABEL=$(echo "$entry" | sed 's/^[^ ]* //')

  dns_query "$IP" "$LABEL" "$DOMAIN" "$TYPE" "$TIMEOUT" &

  ((current_jobs++))
  if [ "$current_jobs" -ge "$MAX_CONCURRENT" ]; then
    # Wait until at least one finishes before starting a new one
    wait -n
    ((current_jobs--))
  fi
done

# Wait for any remaining jobs to complete
wait

# After wait, before counting results
declare -A resolved_values
declare -A answer_counts
first_success=""
consensus_answer=""
max_count=0

# Function to normalize results for comparison
normalize_results() {
  echo "$1" | tr ' ' '\n' | sort | tr '\n' ' ' | sed 's/ $//'
}

# Count results
for entry in "${SERVERS[@]}"; do
  IP=$(echo "$entry" | awk '{print $1}')
  LABEL=$(echo "$entry" | sed 's/^[^ ]* //')
  safe_domain=$(get_safe_filename "$DOMAIN")
  tmp_prefix="/tmp/dns_check_${safe_domain}_${TYPE}_$$"
  
  if [ -f "${tmp_prefix}_${IP}" ]; then
    if [ "$(cat "${tmp_prefix}_${IP}")" = "success" ]; then
      ((successful_queries++))
      if [ -f "${tmp_prefix}_${IP}.resolved" ]; then
        resolved=$(normalize_results "$(cat "${tmp_prefix}_${IP}.resolved")")
        [ -z "$first_success" ] && first_success="$resolved"
        
        # Track answer frequencies
        count="${answer_counts[$resolved]:-0}"
        ((count++))
        answer_counts[$resolved]=$count
        
        # Update consensus if this answer has higher count
        if [ "$count" -gt "$max_count" ]; then
          max_count=$count
          consensus_answer=$resolved
        fi
        
        if [ "$resolved" != "$first_success" ] && [ "$VERBOSE" != "1" ]; then
          printf " ⚠️  %s returned different result:\n   └─ %s\n" "$LABEL" "$resolved"
        fi
      fi
    else
      ((failed_queries++))
    fi
    rm -f "${tmp_prefix}_${IP}" "${tmp_prefix}_${IP}.resolved"
  fi
done

end_time=$(get_time)
duration=$((end_time - script_start_time))
if [ "$VERBOSE" = "1" ]; then
  printf "\nDone! 🏁 🏎️  %s servers in %s ms\n" "$total_servers" "$duration"
  printf "Successful: %d, Failed: %d, Total: %d\n" "$successful_queries" "$failed_queries" "$total_servers"
else
  printf "\n%d/%d succeeded in %dms" "$successful_queries" "$total_servers" "$duration"
  [ "$failed_queries" -gt 0 ] && printf " (%d failed)" "$failed_queries"
  echo
fi

# Function to print sorted answers
print_sorted_answers() {
  local -n answers=$1
  local indent="$2"
  # Create array of "count answer" pairs
  local pairs=()
  for answer in "${!answers[@]}"; do
    pairs+=("${answers[$answer]} $answer")
  done
  # Sort numerically in reverse order
  printf '%s\n' "${pairs[@]}" | sort -rn | while read -r count answer; do
    local percent=$(( (count * 100) / total_servers ))
    printf "${indent}%d%% (%d servers): %s\n" "$percent" "$count" "$answer"
  done
}

# Before exit, show consensus and all answers
if [ ! -z "$consensus_answer" ]; then
  if [ "$VERBOSE" = "1" ]; then
    printf "\n📊 DNS Results:\n"
    print_sorted_answers answer_counts "  "
  else
    consensus_percent=$(( (max_count * 100) / total_servers ))
    printf "✅ Winning Answer: %s (%d%%)\n" "$consensus_answer" "$consensus_percent"
    if [ "${#answer_counts[@]}" -gt 1 ]; then
      printf "📊 Other answers found:\n"
      print_sorted_answers answer_counts "  "
    fi
  fi
fi

[ "$failed_queries" -gt 0 ] && exit 1 || exit 0