#!/usr/bin/env bash
# DNS Geo-Awareness Matrix (with redirection detection)
#
# Inputs:
#   - Domains/IPs:   example.com  www.youtube.com  142.250.190.14
#   - Resolvers:     @8.8.8.8  @1.1.1.1   (if any @IP are provided, ONLY those resolvers are used)
#
# Behavior:
#   - For each domain/IP:
#       * If domain: resolve via each resolver (A records), dedupe by /24, keep up to 3 per resolver.
#       * If an IP literal is provided as the “domain”, use that IP directly (no DNS lookups).
#       * Build a domain-level unique IP set; probe each unique IP **once**:
#           - ICMP ping first; if timeout and hping3 present, TCP:443 fallback.
#           - Best RTT column is 12 chars (left-justified); ICMP: "87ms", TCP: "87ms (http)".
#           - If RTT > 999ms, re-probe once; if still >999ms, show "-" and exclude from stats/GeoHits.
#       * Show the fastest IP per resolver (IP, CC[2], RTT, PTR).
#   - Pre-test: checks for DNS interception/redirect and warns (but continues).
#   - Summary:
#       * Ranks resolvers by average RTT (numeric entries only).
#       * GeoHits (within 25ms of the per-domain best) are shown only if resolver count >= 3.
#
# Requirements:
#   - Required: dig, ping
#   - Optional: whois, geoiplookup (Country/Owner), hping3 (TCP fallback)

set -uo pipefail

# -------------------------- Width detection & formatting --------------------------
WIDTH=80
if [ -t 1 ]; then
  if command -v tput >/dev/null 2>&1; then
    cols="$(tput cols 2>/dev/null || true)"
    [[ "$cols" =~ ^[0-9]+$ ]] && WIDTH="$cols"
  fi
  if ! [[ "$WIDTH" =~ ^[0-9]+$ ]]; then
    cols="$(stty size 2>/dev/null | awk '{print $2}')" || true
    [[ "$cols" =~ ^[0-9]+$ ]] && WIDTH="$cols"
  fi
fi
(( WIDTH < 80 )) && WIDTH=80

FIXED_SUM=0; PIPES=5; PTR_WIDTH=0; RULE_LEN=0; FMT=""

recompute_format() {
  FIXED_SUM=$((16 + 15 + 15 + 2 + 12))  # Resolver + DNS-IP + Answer-IP + CC + BestRTT
  PIPES=5
  PTR_WIDTH=$(( WIDTH - (FIXED_SUM + PIPES) ))
  (( PTR_WIDTH < 8 )) && PTR_WIDTH=8
  FMT="%-16.16s|%-15.15s|%-15.15s|%-2.2s|%-12.12s|%-${PTR_WIDTH}.${PTR_WIDTH}s\n"
  RULE_LEN=$(( FIXED_SUM + PIPES + PTR_WIDTH ))
}
recompute_format

trap '
  if [ -t 1 ] && command -v tput >/dev/null 2>&1; then
    cols="$(tput cols 2>/dev/null || true)"
    [[ "$cols" =~ ^[0-9]+$ ]] && WIDTH="$cols"
    (( WIDTH < 80 )) && WIDTH=80
  fi
  recompute_format
' WINCH

fmtline() { printf "$FMT" "$@"; }
print_rule() { printf -- "%-${RULE_LEN}.${RULE_LEN}s\n" "$(printf "%${RULE_LEN}s" "" | tr " " "-")"; }

print_limited_line() {
  local prefix="$1" text="$2"
  local plen=${#prefix}
  local maxlen=$(( WIDTH - plen ))
  (( maxlen < 8 )) && maxlen=8
  if (( ${#text} > maxlen )); then
    printf "%s%s…\n" "$prefix" "${text:0:maxlen-1}"
  else
    printf "%s%s\n" "$prefix" "$text"
  fi
}

# -------------------------- Parse args --------------------------
DEFAULT_CDN_DOMAINS=( www.cloudflare.com www.youtube.com www.cnn.com www.akamai.com )
EXTRA_RESOLVERS=()
DOMAINS=()

is_ipv4() { [[ "$1" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]; }

if (( $# == 0 )); then
  DOMAINS=( "${DEFAULT_CDN_DOMAINS[@]}" )
else
  for arg in "$@"; do
    if [[ "$arg" == @* ]]; then
      ip="${arg#@}"
      if is_ipv4 "$ip"; then
        EXTRA_RESOLVERS+=( "Custom-$(( ${#EXTRA_RESOLVERS[@]} + 1 )):${ip}" )
      else
        echo "WARN: ignoring invalid resolver token: $arg" >&2
      fi
    else
      DOMAINS+=( "$arg" )
    fi
  done
  # If user passed only @resolvers and no domains/IPs, use defaults
  if (( ${#DOMAINS[@]} == 0 )); then
    DOMAINS=( "${DEFAULT_CDN_DOMAINS[@]}" )
  fi
fi

# -------------------------- Dependencies --------------------------
have_cmd() { command -v "$1" >/dev/null 2>&1; }
req=(dig ping)
missing=()
for c in "${req[@]}"; do have_cmd "$c" || missing+=("$c"); done
if ((${#missing[@]})); then
  echo "ERROR: Missing required command(s): ${missing[*]}" >&2
  exit 1
fi
have_whois()   { have_cmd whois; }
have_geoip()   { have_cmd geoiplookup; }
have_hping3()  { have_cmd hping3;    }

# -------------------------- Build resolver list --------------------------
if (( ${#EXTRA_RESOLVERS[@]} > 0 )); then
  RESOLVERS=( "${EXTRA_RESOLVERS[@]}" )
else
  RESOLVERS=(
    "Google:8.8.8.8"
    "Cloudflare:1.1.1.1"
    "Quad9:9.9.9.9"
    "Quad9-Unfiltered:9.9.9.10"
    "OpenDNS:208.67.222.222"
    "AdGuard:94.140.14.14"
    "CleanBrowsing:185.228.168.9"
    "Neustar:156.154.70.1"
    "Level3:4.2.2.2"
  )
  SYS_DNS="$(awk '/^nameserver/{print $2}' /etc/resolv.conf | grep -v -E '^(127\.0\.0\.1|::1)$' | head -n1 || true)"
  if [[ -n "${SYS_DNS:-}" ]]; then
    RESOLVERS+=( "SystemDefault:${SYS_DNS}" )
  fi
fi

# Pretty list of resolvers for the header
join_resolvers() {
  local s="" name ip
  for entry in "${RESOLVERS[@]}"; do
    name="${entry%%:*}"; ip="${entry#*:}"
    [[ -n "$s" ]] && s+=", "
    s+="${name}(${ip})"
  done
  printf "%s" "$s"
}

echo
echo "DNS Geo-Awareness Matrix (with redirection detection)"
print_limited_line "Domains under test: " "${DOMAINS[*]}"
print_limited_line "Resolvers under test: " "$(join_resolvers)"
echo

# -------------------------- Helpers (RAM caches only) --------------------------
# Per-run in-memory caches (no disk writes)
declare -A PTR_CACHE COUNTRY_CACHE OWNER_CACHE

resolve_ips_v4() {
  local resolver_ip="$1" name_or_ip="$2"
  if is_ipv4 "$name_or_ip"; then
    printf "%s\n" "$name_or_ip"
    return
  fi
  dig +short A "$name_or_ip" @"$resolver_ip" 2>/dev/null | grep -E '^[0-9]+\.' || true
}

dedupe_by24_and_limit3() { awk -F. '!seen[$1"."$2"."$3]++' | head -n 3; }

ptr_name() {
  local ip="$1"
  local cached="${PTR_CACHE[$ip]:-}"
  if [[ -n "$cached" ]]; then echo "$cached"; return; fi
  # Use dig defaults (time=5, tries=3)
  local p
  p="$(dig +short -x "$ip" 2>/dev/null | sed 's/\.$//' | head -n1)"
  [[ -z "$p" ]] && p="-"
  PTR_CACHE["$ip"]="$p"
  echo "$p"
}

get_country(){
  local ip="$1"
  local cached="${COUNTRY_CACHE[$ip]:-}"
  if [[ -n "$cached" ]]; then echo "$cached"; return; fi
  local c=""
  if have_whois; then
    c="$(timeout 5 whois -H "$ip" 2>/dev/null | awk 'BEGIN{IGNORECASE=1}/^country[[:space:]]*:/ {print $2; exit}')" || true
  fi
  if [[ -z "$c" ]] && have_geoip; then
    c="$(geoiplookup "$ip" 2>/dev/null | sed -n '1s/.*: //p')" || true
  fi
  [[ -z "$c" ]] && c="UNKNOWN"
  COUNTRY_CACHE["$ip"]="$c"
  echo "$c"
}

get_owner(){
  local ip="$1"
  local cached="${OWNER_CACHE[$ip]:-}"
  if [[ -n "$cached" ]]; then echo "$cached"; return; fi
  if ! have_whois; then echo "UNKNOWN"; return; fi
  local o
  o="$(timeout 5 whois -H "$ip" 2>/dev/null | awk '
      BEGIN{IGNORECASE=1}
      /^(OrgName|org-name|owner|organisation|descr|netname)[[:space:]]*:/ {sub(/^[^:]+:[[:space:]]*/,""); print; exit}
    ' | sed -E 's/[[:space:]]+$//' | head -n1)" || true
  [[ -z "$o" ]] && o="UNKNOWN"
  OWNER_CACHE["$ip"]="$o"
  echo "$o"
}

get_icmp_avg_ms(){
  local ip="$1" avg=""
  if timeout 4 ping -c 3 -i 0.2 -W 1 "$ip" >/tmp/ping.$$ 2>/dev/null; then
    avg="$(awk -F'/' '/^rtt/ {print $(NF-2)}' /tmp/ping.$$)" || true
    [[ -z "$avg" ]] && avg="$(awk -F'/' '/round-trip/ {print $(NF-2)}' /tmp/ping.$$)" || true
  fi
  rm -f /tmp/ping.$$ || true
  [[ -n "$avg" ]] && echo "$avg" || echo "timeout"
}

get_tcp443_avg_ms(){
  local ip="$1"
  if ! have_hping3; then echo "no-tcp"; return; fi
  # Accept both "time=X ms" and "rtt=X ms" variants from hping3
  local out
  out="$(timeout 8 hping3 -S -p 443 -c 3 "$ip" 2>/dev/null \
         | grep -oE '(time|rtt)=[0-9.]+ ms' \
         | grep -oE '[0-9.]+' )"
  [[ -z "$out" ]] && { echo "timeout"; return; }
  awk 'BEGIN{sum=0; n=0} {sum+=$1; n++} END{if(n>0) printf("%.3f", sum/n); else print "timeout"}' <<<"$out"
}

# Re-probe logic for >999ms (return: "-" if still too high; else integer ms)
reprobe_if_too_high(){
  local ip="$1" method="$2" val="$3"
  if [[ "$val" =~ ^[0-9]+(\.[0-9]+)?$ ]]; then
    awk -v x="$val" 'BEGIN{exit (x>999)?0:1}'; hi=$?
    if (( hi == 0 )); then
      local again=""
      if [[ "$method" == "ICMP" ]]; then
        again="$(get_icmp_avg_ms "$ip")"
      else
        again="$(get_tcp443_avg_ms "$ip")"
      fi
      if [[ "$again" =~ ^[0-9]+(\.[0-9]+)?$ ]]; then
        awk -v y="$again" 'BEGIN{exit (y>999)?0:1}'; hi2=$?
        if (( hi2 == 0 )); then
          echo "-"
        else
          printf "%.0f" "$again"
        fi
      else
        echo "-"
      fi
    else
      printf "%.0f" "$val"
    fi
  else
    echo "$val"
  fi
}

# Probe one IP -> "METHOD|VALUE"
probe_ip(){
  local ip="$1"
  local icmp tcp
  icmp="$(get_icmp_avg_ms "$ip")"
  if [[ "$icmp" != "timeout" ]]; then
    local adj="$(reprobe_if_too_high "$ip" "ICMP" "$icmp")"
    echo "ICMP|$adj"; return 0
  fi
  tcp="$(get_tcp443_avg_ms "$ip")"
  if [[ "$tcp" != "timeout" && "$tcp" != "no-tcp" ]]; then
    local adj="$(reprobe_if_too_high "$ip" "TCP443" "$tcp")"
    echo "TCP443|$adj"; return 0
  fi
  echo "NONE|$tcp"   # "timeout" or "no-tcp"
  return 0
}

# -------------------------- Redirection checks (pre-test) --------------------------
warn_flags=()

# 1) Non-DNS probe (should not answer)
bogus_ip="203.0.113.1"
out="$(dig @${bogus_ip} example.com A +time=2 +tries=1 +stats 2>/dev/null || true)"
if grep -q "^;; SERVER:" <<<"$out"; then
  actual="$(sed -n 's/^;; SERVER: \([^#]*\).*/\1/p' <<<"$out" | head -n1 | tr -d ' ')"
  warn_flags+=("Non-DNS probe on ${bogus_ip}:53 returned a DNS response (responder: ${actual})")
fi

# 2) Identity checks per resolver (UDP & TCP)
declare -A redir_udp redir_tcp redir_to
check_identity(){
  local name="$1" intended="$2" proto="$3" ident actual
  if [[ "$proto" == "udp" ]]; then
    ident="$(dig @${intended} example.com A +time=2 +tries=1 +stats 2>/dev/null || true)"
  else
    ident="$(dig @${intended} example.com A +tcp +time=2 +tries=1 +stats 2>/dev/null || true)"
  fi
  if grep -q "^;; SERVER:" <<<"$ident"; then
    actual="$(sed -n 's/^;; SERVER: \([^#]*\).*/\1/p' <<<"$ident" | head -n1 | tr -d ' ')"
    if [[ -n "$actual" && "$actual" != "$intended" ]]; then
      warn_flags+=("${name} (${intended}) over ${proto^^} appears redirected to ${actual}")
      if [[ "$proto" == "udp" ]]; then redir_udp["$name"]=1; redir_to["$name,$proto"]="$actual"; else redir_tcp["$name"]=1; redir_to["$name,$proto"]="$actual"; fi
    fi
  fi
}

for entry in "${RESOLVERS[@]}"; do
  name="${entry%%:*}"; dnsip="${entry#*:}"
  check_identity "$name" "$dnsip" udp
  check_identity "$name" "$dnsip" tcp
done

if ((${#warn_flags[@]})); then
  echo "⚠ Detected possible DNS interception/redirect BEFORE tests:"
  for w in "${warn_flags[@]}"; do echo "  - $w"; done
  echo "⚠ Results below may be tainted due to DNS redirect."
else
  echo "✓ No interception/redirect indications detected by the simple checks."
fi
echo

# -------------------------- Main matrix --------------------------
declare -A rtt_sum rtt_count geo_hits

for domain in "${DOMAINS[@]}"; do
  echo "=== ${domain} ==="
  fmtline "Resolver" "DNS-IP" "Answer-IP" "CC" "Best RTT" "PTR (reverse DNS)"
  print_rule

  declare -A RES_IPS SEEN_IP IP_RESULT
  DOMAIN_UNIQUE_IPS=()

  # Phase 1: resolve all resolvers; build unique set
  for entry in "${RESOLVERS[@]}"; do
    name="${entry%%:*}"; dnsip="${entry#*:}"
    ips_raw="$(resolve_ips_v4 "$dnsip" "$domain" || true)"
    ips="$(sed '/^$/d' <<<"$ips_raw" | dedupe_by24_and_limit3 || true)"
    RES_IPS["$name"]="$ips"
    while IFS= read -r ip; do
      [[ -z "$ip" ]] && continue
      if [[ -z "${SEEN_IP["$ip"]:-}" ]]; then
        SEEN_IP["$ip"]=1
        DOMAIN_UNIQUE_IPS+=("$ip")
      fi
    done <<< "$ips"
  done

  # Phase 2: probe each unique IP once
  for ip in "${DOMAIN_UNIQUE_IPS[@]}"; do
    IP_RESULT["$ip"]="$(probe_ip "$ip")"
  done

  # Phase 3: per-resolver best, print rows, track domain best
  domain_best=""
  declare -A fast_rtt_by_resolver

  for entry in "${RESOLVERS[@]}"; do
    name="${entry%%:*}"; dnsip="${entry#*:}"; ips="${RES_IPS["$name"]}"

    if [[ -z "$ips" ]]; then
      fmtline "$name" "$dnsip" "<no answer>" "--" "timeout" "-"
      continue
    fi

    best_ip=""; best_val=""; best_method=""
    while IFS= read -r ip; do
      [[ -z "$ip" ]] && continue
      method="${IP_RESULT["$ip"]%%|*}"
      value="${IP_RESULT["$ip"]#*|}"   # ms number, "-", timeout, no-tcp
      if [[ "$value" =~ ^[0-9]+(\.[0-9]+)?$ ]]; then
        if [[ -z "$best_val" || $(awk -v a="$value" -v b="$best_val" 'BEGIN{print (a<b)?1:0}') -eq 1 ]]; then
          best_val="$value"; best_ip="$ip"; best_method="$method"
        fi
      fi
    done <<< "$ips"

    if [[ -n "$best_ip" ]]; then
      cc="$(get_country "$best_ip")"; cc="${cc:0:2}"
      ptr="$(ptr_name "$best_ip")"
      rtt_disp="$(printf "%.0f" "$best_val")ms"
      [[ "$best_method" == "TCP443" ]] && rtt_disp="${rtt_disp} (http)"
      fmtline "$name" "$dnsip" "$best_ip" "$cc" "$rtt_disp" "$ptr"

      rtt_int="${best_val%.*}"
      rtt_sum["$name"]=$(( ${rtt_sum["$name"]:-0} + rtt_int ))
      rtt_count["$name"]=$(( ${rtt_count["$name"]:-0} + 1 ))
      fast_rtt_by_resolver["$name"]="$rtt_int"
      if [[ -z "$domain_best" || $rtt_int -lt $domain_best ]]; then domain_best="$rtt_int"; fi
    else
      one_ip="$(head -n1 <<<"$ips")"
      cc="$(get_country "$one_ip")"; cc="${cc:0:2}"
      ptr="$(ptr_name "$one_ip")"
      status="${IP_RESULT["$one_ip"]#*|}"   # "-", timeout, no-tcp
      fmtline "$name" "$dnsip" "$one_ip" "$cc" "$status" "$ptr"
    fi
  done

  # Geo-hits (only if we have a domain best)
  if [[ -n "$domain_best" ]]; then
    for entry in "${RESOLVERS[@]}"; do
      name="${entry%%:*}"
      val="${fast_rtt_by_resolver["$name"]:-}"
      [[ -z "$val" ]] && continue
      diff=$(( val - domain_best )); (( diff<0 )) && diff=$(( -diff ))
      (( diff <= 25 )) && geo_hits["$name"]=$(( ${geo_hits["$name"]:-0} + 1 ))
    done
  fi

  echo
done

# -------------------------- Summary --------------------------
echo "=== Interception / Redirection Summary ==="
if ((${#warn_flags[@]})); then
  for w in "${warn_flags[@]}"; do
    if [[ "$w" =~ ([0-9]+\.[0-9]+\.[0-9]+\.[0-9]+)$ ]]; then
      ip="${BASH_REMATCH[1]}"
      owner="$(get_owner "$ip")"; country="$(get_country "$ip")"
      echo "  - ${w}  [owner: ${owner:-UNKNOWN}; country: ${country:-UNKNOWN}]"
    else
      echo "  - ${w}"
    fi
  done
else
  echo "No interception/redirect indications detected by the simple checks."
fi
echo

echo "=== Geo-Awareness Ranking (lower avg RTT & more hits is better) ==="
resolver_count="${#RESOLVERS[@]}"
# Build summary format dynamically to fill WIDTH; Notes column flexes
SUM_FIXED=$((16 + 15 + 12))  # Resolver|DNS-IP|AvgRTT
SUM_PIPES=3
if (( resolver_count >= 3 )); then
  SUM_FIXED=$((SUM_FIXED + 8))   # GeoHits (8 wide)
  SUM_PIPES=$((SUM_PIPES + 1))
fi
NOTES_WIDTH=$(( WIDTH - (SUM_FIXED + SUM_PIPES) ))
(( NOTES_WIDTH < 8 )) && NOTES_WIDTH=8

if (( resolver_count >= 3 )); then
  SUMFMT="%-16.16s|%-15.15s|%-12.12s|%-8.8s|%-${NOTES_WIDTH}.${NOTES_WIDTH}s\n"
  printf "$SUMFMT" "Resolver" "DNS-IP" "AvgRTT" "GeoHits" "Notes"
else
  SUMFMT="%-16.16s|%-15.15s|%-12.12s|%-${NOTES_WIDTH}.${NOTES_WIDTH}s\n"
  printf "$SUMFMT" "Resolver" "DNS-IP" "AvgRTT" "Notes"
fi
printf -- "%-${WIDTH}.${WIDTH}s\n" "$(printf "%${WIDTH}s" "" | tr ' ' '-')"

for entry in "${RESOLVERS[@]}"; do
  name="${entry%%:*}"; dnsip="${entry#*:}"
  count=${rtt_count["$name"]:-0}; sum=${rtt_sum["$name"]:-0}
  avg="n/a"; ((count>0)) && avg=$(( sum / count ))
  note=""
  (( ${redir_udp["$name"]:-0} == 1 )) && note+="UDP-redirect "
  (( ${redir_tcp["$name"]:-0} == 1 )) && note+="TCP-redirect "

  if (( resolver_count >= 3 )); then
    hits=${geo_hits["$name"]:-0}
    printf "$SUMFMT" "$name" "$dnsip" "$avg" "$hits" "$note"
  else
    printf "$SUMFMT" "$name" "$dnsip" "$avg" "$note"
  fi
done

if ((${#warn_flags[@]})); then
  echo
  echo "⚠ Results may be tainted due to DNS redirect."
fi

echo
echo "Done."

