#!/usr/bin/env bash

#########################################################################################
#### dns-propagation-checker.sh — Check DNS propagation across public resolvers      ####
#### Queries Cloudflare, Google, Quad9, OpenDNS, compares results                    ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./dns-propagation-checker.sh example.com                                      ####
####   ./dns-propagation-checker.sh example.com --type MX                            ####
####   ./dns-propagation-checker.sh example.com --watch 30                           ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
DOMAIN=""
RECORD_TYPE="A"
TIMEOUT=5
COLOR="auto"
JSON_OUTPUT="false"
WATCH_INTERVAL=0
EXPECTED=""

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME

# ── Built-in Resolvers ───────────────────────────────────────────────
RESOLVER_NAMES=("Cloudflare" "Google" "Quad9" "OpenDNS" "Cloudflare-2" "Google-2")
RESOLVER_IPS=("1.1.1.1" "8.8.8.8" "9.9.9.9" "208.67.222.222" "1.0.0.1" "8.8.4.4")
CUSTOM_RESOLVERS=()

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BOLD="" DIM="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BOLD="" DIM="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
err()  { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
warn() { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }

# ══════════════════════════════════════════════════════════════════════
#                       USAGE
# ══════════════════════════════════════════════════════════════════════

usage() {
    cat <<EOF
${SCRIPT_NAME} — Check DNS propagation across public resolvers

USAGE:
  ${SCRIPT_NAME} DOMAIN [OPTIONS]

OPTIONS:
  --type TYPE           Record type: A, AAAA, MX, CNAME, TXT, NS, SOA, PTR (default: A)
  --resolver IP         Add a custom resolver (repeatable)
  --expected VALUE      Expected answer to verify against
  --json                Output results in JSON format
  --watch SECONDS       Re-check every N seconds until all resolvers agree
  --timeout SECONDS     Per-query timeout in seconds (default: 5)
  --no-color            Disable colored output
  --help                Show this help

EXAMPLES:
  # Check A record propagation
  ./${SCRIPT_NAME} example.com

  # Check MX records
  ./${SCRIPT_NAME} example.com --type MX

  # Watch until propagated
  ./${SCRIPT_NAME} example.com --type A --watch 30

  # Verify specific value
  ./${SCRIPT_NAME} example.com --expected "203.0.113.50"

  # Add custom resolvers
  ./${SCRIPT_NAME} example.com --resolver 4.2.2.1 --resolver 77.88.8.8

  # JSON output
  ./${SCRIPT_NAME} example.com --json
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                       ARGUMENT PARSING
# ══════════════════════════════════════════════════════════════════════

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --type)
                RECORD_TYPE="${2^^}"; shift 2 ;;
            --resolver)
                CUSTOM_RESOLVERS+=("$2"); shift 2 ;;
            --expected)
                EXPECTED="$2"; shift 2 ;;
            --json)
                JSON_OUTPUT="true"; shift ;;
            --watch)
                WATCH_INTERVAL="$2"; shift 2 ;;
            --timeout)
                TIMEOUT="$2"; shift 2 ;;
            --no-color)
                COLOR="never"; shift ;;
            --help|-h)
                setup_colors
                usage
                exit 0 ;;
            -*)
                err "Unknown option: $1"
                echo "Run ${SCRIPT_NAME} --help for usage" >&2
                exit 1 ;;
            *)
                if [[ -z "$DOMAIN" ]]; then
                    DOMAIN="$1"
                else
                    err "Unexpected argument: $1"
                    exit 1
                fi
                shift ;;
        esac
    done

    if [[ -z "$DOMAIN" ]]; then
        err "Domain name is required"
        echo "Run ${SCRIPT_NAME} --help for usage" >&2
        exit 1
    fi

    local valid_types="A AAAA MX CNAME TXT NS SOA PTR"
    if [[ ! " $valid_types " =~ " $RECORD_TYPE " ]]; then
        err "Invalid record type: $RECORD_TYPE"
        exit 1
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       DNS QUERY
# ══════════════════════════════════════════════════════════════════════

query_resolver() {
    local resolver_ip="$1" domain="$2" rtype="$3" timeout="$4"
    local output ttl_output answer ttl

    output=$(dig +time="$timeout" +tries=1 +short "@${resolver_ip}" "$domain" "$rtype" 2>/dev/null) || true
    ttl_output=$(dig +time="$timeout" +tries=1 +noall +answer "@${resolver_ip}" "$domain" "$rtype" 2>/dev/null) || true
    answer=$(echo "$output" | tr '\n' ' ' | sed 's/ *$//')
    ttl=$(echo "$ttl_output" | awk 'NR==1{print $2}')

    if [[ -z "$answer" ]]; then
        echo "FAIL||"; return
    fi
    echo "${answer}|${ttl:-?}|"
}

# ══════════════════════════════════════════════════════════════════════
#                       MAJORITY ANSWER
# ══════════════════════════════════════════════════════════════════════

find_majority() {
    local -n answers_ref=$1
    local -A counts
    local max_count=0 majority=""
    for answer in "${answers_ref[@]}"; do
        [[ "$answer" == "FAIL" ]] && continue
        counts["$answer"]=$(( ${counts["$answer"]:-0} + 1 ))
        if [[ ${counts["$answer"]} -gt $max_count ]]; then
            max_count=${counts["$answer"]}; majority="$answer"
        fi
    done
    echo "$majority"
}

# ══════════════════════════════════════════════════════════════════════
#                       RUN CHECK
# ══════════════════════════════════════════════════════════════════════

run_check() {
    local all_names=("${RESOLVER_NAMES[@]}")
    local all_ips=("${RESOLVER_IPS[@]}")

    for custom in "${CUSTOM_RESOLVERS[@]}"; do
        all_names+=("Custom-${custom}")
        all_ips+=("$custom")
    done

    local total=${#all_names[@]}
    local answers=() ttls=()

    for i in $(seq 0 $(( total - 1 ))); do
        local result
        result=$(query_resolver "${all_ips[$i]}" "$DOMAIN" "$RECORD_TYPE" "$TIMEOUT")
        answers+=("$(echo "$result" | cut -d'|' -f1)")
        ttls+=("$(echo "$result" | cut -d'|' -f2)")
    done

    local majority
    majority=$(find_majority answers)
    local compare_to="${EXPECTED:-$majority}"
    local agree_count=0 statuses=()

    for i in $(seq 0 $(( total - 1 ))); do
        if [[ "${answers[$i]}" == "FAIL" ]]; then
            statuses+=("FAIL")
        elif [[ "${answers[$i]}" == "$compare_to" ]]; then
            statuses+=("MATCH"); agree_count=$((agree_count + 1))
        else
            statuses+=("MISMATCH")
        fi
    done

    if [[ "$JSON_OUTPUT" == "true" ]]; then
        print_json all_names all_ips answers ttls statuses "$agree_count" "$total" "$majority"
    else
        print_table all_names all_ips answers ttls statuses "$agree_count" "$total" "$majority" "$compare_to"
    fi
    [[ "$agree_count" -eq "$total" ]] && return 0 || return 1
}

# ══════════════════════════════════════════════════════════════════════
#                       OUTPUT: TABLE
# ══════════════════════════════════════════════════════════════════════

print_table() {
    local -n names_ref=$1 ips_ref=$2 ans_ref=$3 ttl_ref=$4 stat_ref=$5
    local agree="$6" total="$7" majority="$8" compare="$9"

    echo ""
    echo -e "${BOLD}DNS Propagation Check — ${DOMAIN} (${RECORD_TYPE})${RESET}"
    echo -e "${DIM}$(date -u '+%Y-%m-%d %H:%M:%S UTC')${RESET}"
    echo ""
    printf "  ${BOLD}%-20s %-17s %-22s %-6s %s${RESET}\n" "RESOLVER" "IP" "RESULT" "TTL" "STATUS"
    printf "  %s\n" "$(printf '%.0s─' {1..78})"

    local count=${#names_ref[@]}
    for i in $(seq 0 $(( count - 1 ))); do
        local color status_str
        case "${stat_ref[$i]}" in
            MATCH)    color="$GREEN";  status_str="MATCH" ;;
            MISMATCH) color="$YELLOW"; status_str="MISMATCH" ;;
            FAIL)     color="$RED";    status_str="FAIL" ;;
        esac

        local display_answer="${ans_ref[$i]}"
        if [[ ${#display_answer} -gt 20 ]]; then
            display_answer="${display_answer:0:17}..."
        fi

        printf "  %-20s %-17s %b%-22s%b %-6s %b%s%b\n" \
            "${names_ref[$i]}" \
            "${ips_ref[$i]}" \
            "$color" "$display_answer" "$RESET" \
            "${ttl_ref[$i]}" \
            "$color" "$status_str" "$RESET"
    done

    echo ""
    echo -e "  ${BOLD}Summary${RESET}"
    if [[ -n "$EXPECTED" ]]; then
        printf "  %-20s %s\n" "Expected answer:" "$EXPECTED"
    fi
    printf "  %-20s %s\n" "Majority answer:" "${majority:-N/A}"
    printf "  %-20s %s\n" "Agree:" "${agree}/${total} resolvers"

    if [[ "$agree" -eq "$total" ]]; then
        printf "  %-20s " "Status:"; echo -e "${GREEN}PROPAGATION COMPLETE${RESET}"
    else
        printf "  %-20s " "Status:"; echo -e "${YELLOW}PROPAGATION PENDING${RESET}"
    fi
    echo ""
}

# ══════════════════════════════════════════════════════════════════════
#                       OUTPUT: JSON
# ══════════════════════════════════════════════════════════════════════

print_json() {
    local -n jnames=$1 jips=$2 jans=$3 jttls=$4 jstats=$5
    local agree="$6" total="$7" majority="$8"
    local count=${#jnames[@]} propagated="false"
    [[ "$agree" -eq "$total" ]] && propagated="true"

    printf '{"domain":"%s","type":"%s","timestamp":"%s","results":[' \
        "$DOMAIN" "$RECORD_TYPE" "$(date -u '+%Y-%m-%dT%H:%M:%SZ')"
    for i in $(seq 0 $(( count - 1 ))); do
        [[ $i -gt 0 ]] && printf ','
        local escaped_answer
        escaped_answer=$(echo "${jans[$i]}" | sed 's/"/\\"/g')
        printf '{"resolver":"%s","ip":"%s","answer":"%s","ttl":"%s","status":"%s"}' \
            "${jnames[$i]}" "${jips[$i]}" "$escaped_answer" "${jttls[$i]}" "${jstats[$i]}"
    done
    local escaped_majority
    escaped_majority=$(echo "$majority" | sed 's/"/\\"/g')
    printf '],"summary":{"majority":"%s","agree":%d,"total":%d,"propagated":%s}}\n' \
        "$escaped_majority" "$agree" "$total" "$propagated"
}

# ══════════════════════════════════════════════════════════════════════
#                       MAIN
# ══════════════════════════════════════════════════════════════════════

main() {
    parse_args "$@"
    setup_colors

    if ! command -v dig &>/dev/null; then
        err "dig is required but not found. Install dnsutils (Debian/Ubuntu) or bind-utils (RHEL/CentOS)."
        exit 1
    fi

    if [[ "$WATCH_INTERVAL" -gt 0 ]]; then
        local cycle=1
        while true; do
            if [[ "$JSON_OUTPUT" != "true" ]]; then
                [[ $cycle -gt 1 ]] && echo -e "${DIM}────────────────────────────────────────────────${RESET}"
                echo -e "${DIM}Watch cycle ${cycle} — checking every ${WATCH_INTERVAL}s (Ctrl+C to stop)${RESET}"
            fi
            if run_check; then
                [[ "$JSON_OUTPUT" != "true" ]] && echo -e "  ${GREEN}All resolvers agree. Propagation complete.${RESET}\n"
                exit 0
            fi
            cycle=$((cycle + 1))
            sleep "$WATCH_INTERVAL"
        done
    else
        run_check && exit 0 || exit 1
    fi
}

main "$@"