#!/usr/bin/env bash

#####################################################################################
#### dns-smoke-tests.sh — Verify DNS infrastructure is healthy                   ####
#### Checks resolution, zone transfers, SOA, DNSSEC, response time, DoT.        ####
####                                                                             ####
#### Author:  Phil Connor                                                        ####
#### Contact: contact@mylinux.work                                               ####
#### License: MIT                                                                ####
#### Version: 1.0                                                                ####
####                                                                             ####
#### Usage:   ./dns-smoke-tests.sh                                               ####
####          DNS_SERVER=192.168.1.1 DOMAIN=example.com ./dns-smoke-tests.sh     ####
####                                                                             ####
#### See --help for all options.                                                 ####
#####################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
DNS_SERVER="${DNS_SERVER:-}"
DOMAIN="${DOMAIN:-example.com}"
REVERSE_IP="${REVERSE_IP:-}"
ZONE="${ZONE:-}"
ZONE_MASTER="${ZONE_MASTER:-}"
DNSSEC_DOMAIN="${DNSSEC_DOMAIN:-}"
DOT_SERVER="${DOT_SERVER:-}"
MAX_RESPONSE_MS="${MAX_RESPONSE_MS:-500}"
TEST_RECORDS="${TEST_RECORDS:-}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
PASS=0; FAIL=0; SKIP=0; TOTAL=0
RESULTS=()
START_TIME=""

# ── Dig tool detection ───────────────────────────────────────────────
DIG_CMD=""
detect_dig() {
    if command -v dig >/dev/null 2>&1; then
        DIG_CMD="dig"
    elif command -v drill >/dev/null 2>&1; then
        DIG_CMD="drill"
    else
        err "Neither dig nor drill found. Install dnsutils or ldns."
        exit 1
    fi
    verbose "Using ${DIG_CMD} for DNS queries"
}

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

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${BLUE}[DEBUG]${RESET} $*"; fi; }

# ── Test Result Recording ─────────────────────────────────────────────
record_pass() {
    local name="$1" detail="${2:-}"
    ((PASS++)) || true; ((TOTAL++)) || true
    RESULTS+=("PASS|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then echo "ok ${TOTAL} - ${name}${detail:+ (${detail})}"
    else echo -e "  ${GREEN}✓${RESET} ${name}${detail:+ — ${detail}}"; fi
}

record_fail() {
    local name="$1" detail="${2:-}"
    ((FAIL++)) || true; ((TOTAL++)) || true
    RESULTS+=("FAIL|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "not ok ${TOTAL} - ${name}"
        [[ -n "$detail" ]] && echo "  # ${detail}"
    else echo -e "  ${RED}✗${RESET} ${name}${detail:+ — ${detail}}"; fi
}

record_skip() {
    local name="$1" reason="${2:-}"
    ((SKIP++)) || true; ((TOTAL++)) || true
    RESULTS+=("SKIP|${name}|${reason}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then echo "ok ${TOTAL} - ${name} # SKIP ${reason}"
    else echo -e "  ${YELLOW}⊘${RESET} ${name}${reason:+ — ${reason}}"; fi
}

# ── Helpers ───────────────────────────────────────────────────────────
has_cmd() { command -v "$1" >/dev/null 2>&1; }

# Build dig command with optional @server
dig_cmd() {
    if [[ -n "$DNS_SERVER" ]]; then
        "$DIG_CMD" "@${DNS_SERVER}" "$@"
    else
        "$DIG_CMD" "$@"
    fi
}

# ── Output Functions ──────────────────────────────────────────────────
section_header() {
    local name="$1"
    if [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo ""
        echo -e "${BOLD}${name}${RESET}"
    fi
}

print_header() {
    if [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo ""
        echo -e "${BOLD}DNS Smoke Tests${RESET}"
        echo "Domain:     ${DOMAIN}"
        [[ -n "$DNS_SERVER" ]] && echo "Server:     ${DNS_SERVER}" || echo "Server:     (system resolver)"
        echo "Time:       $(date -u +%Y-%m-%dT%H:%M:%SZ)"
    fi
}

print_tap_header() {
    echo "TAP version 13"
}

print_summary() {
    local end_time; end_time=$(date +%s)
    local duration=$(( end_time - START_TIME ))
    echo ""
    echo -e "${BOLD}────────────────────────────────────────${RESET}"
    echo -e "${BOLD}Summary${RESET}  ${DOMAIN}  ${DNS_SERVER:-(system resolver)}"
    echo -e "  ${GREEN}${PASS} passed${RESET}  ${RED}${FAIL} failed${RESET}  ${YELLOW}${SKIP} skipped${RESET}  (${duration}s)"
    echo -e "${BOLD}────────────────────────────────────────${RESET}"
    if [[ $FAIL -eq 0 ]]; then echo -e "${GREEN}${BOLD}All tests passed.${RESET}"
    else echo -e "${RED}${BOLD}${FAIL} test(s) failed.${RESET}"; fi
}

print_tap_footer() {
    echo "1..${TOTAL}"
    echo "# pass ${PASS}"
    echo "# fail ${FAIL}"
    echo "# skip ${SKIP}"
}

# ══════════════════════════════════════════════════════════════════════
#                          TESTS
# ══════════════════════════════════════════════════════════════════════

# ── 1. Resolver Reachable ─────────────────────────────────────────────
test_resolver_reachable() {
    section_header "Connectivity"
    local output
    output=$(dig_cmd +short +time=5 +tries=1 "${DOMAIN}" A 2>&1) || true
    if [[ -n "$output" ]] && ! echo "$output" | grep -qi "timed out\|connection refused\|no servers\|SERVFAIL"; then
        record_pass "Resolver reachable" "${DNS_SERVER:-(system resolver)}"
    else
        record_fail "Resolver reachable" "${DNS_SERVER:-(system resolver)} — ${output:-no response}"
    fi
}

# ── 2. Forward Resolution (A) ────────────────────────────────────────
test_forward_resolution() {
    section_header "Resolution"
    local output
    output=$(dig_cmd +short "${DOMAIN}" A 2>/dev/null) || true
    if [[ -n "$output" ]]; then
        local first_ip
        first_ip=$(echo "$output" | head -1)
        record_pass "Forward resolution (${DOMAIN} A)" "${first_ip}"
    else
        record_fail "Forward resolution (${DOMAIN} A)" "no A record returned"
    fi
}

# ── 3. AAAA Resolution ───────────────────────────────────────────────
test_aaaa_resolution() {
    local output
    output=$(dig_cmd +short "${DOMAIN}" AAAA 2>/dev/null) || true
    if [[ -n "$output" ]]; then
        local first_ip
        first_ip=$(echo "$output" | head -1)
        record_pass "AAAA resolution (${DOMAIN})" "${first_ip}"
    else
        record_skip "AAAA resolution (${DOMAIN})" "no AAAA record"
    fi
}

# ── 4. MX Resolution ─────────────────────────────────────────────────
test_mx_resolution() {
    local output
    output=$(dig_cmd +short "${DOMAIN}" MX 2>/dev/null) || true
    if [[ -n "$output" ]]; then
        local first_mx
        first_mx=$(echo "$output" | head -1)
        record_pass "MX resolution (${DOMAIN})" "${first_mx}"
    else
        record_skip "MX resolution (${DOMAIN})" "no MX record"
    fi
}

# ── 5. Reverse Lookup ────────────────────────────────────────────────
test_reverse_lookup() {
    if [[ -z "$REVERSE_IP" ]]; then
        record_skip "Reverse lookup" "REVERSE_IP not set"
        return
    fi
    local output
    output=$(dig_cmd +short -x "${REVERSE_IP}" 2>/dev/null) || true
    if [[ -n "$output" ]]; then
        record_pass "Reverse lookup (${REVERSE_IP})" "${output}"
    else
        record_fail "Reverse lookup (${REVERSE_IP})" "no PTR record returned"
    fi
}

# ── 6. Response Time ─────────────────────────────────────────────────
test_response_time() {
    section_header "Performance"
    local output query_time
    output=$(dig_cmd "${DOMAIN}" A 2>/dev/null) || true
    # dig outputs "Query time: 12 msec" or ";; Query time: 12 msec"
    query_time=$(echo "$output" | grep -i "query time" | grep -oP '[0-9]+' | head -1) || true
    if [[ -z "$query_time" ]]; then
        # drill outputs ";; Query time: 0 msec"
        query_time=$(echo "$output" | grep -i "query time" | awk '{print $4}') || true
    fi
    if [[ -n "$query_time" ]]; then
        if [[ "$query_time" -le "$MAX_RESPONSE_MS" ]]; then
            record_pass "Response time" "${query_time}ms (<= ${MAX_RESPONSE_MS}ms)"
        else
            record_fail "Response time" "${query_time}ms (> ${MAX_RESPONSE_MS}ms)"
        fi
    else
        record_fail "Response time" "could not parse query time"
    fi
}

# ── 7. Authoritative Answer ──────────────────────────────────────────
test_authoritative_answer() {
    section_header "Authority"
    local output
    output=$(dig_cmd "${DOMAIN}" A 2>/dev/null) || true
    if echo "$output" | grep -q "flags:.*aa"; then
        record_pass "Authoritative answer (${DOMAIN})" "AA flag set"
    else
        record_fail "Authoritative answer (${DOMAIN})" "AA flag not set — server is not authoritative"
    fi
}

# ── 8. SOA Serial ────────────────────────────────────────────────────
test_soa_serial() {
    local output serial
    output=$(dig_cmd +short "${DOMAIN}" SOA 2>/dev/null) || true
    if [[ -z "$output" ]]; then
        record_fail "SOA serial (${DOMAIN})" "no SOA record returned"
        return
    fi
    # SOA format: ns1.example.com. admin.example.com. 2026051201 3600 900 604800 86400
    serial=$(echo "$output" | awk '{print $3}') || true
    if [[ -z "$serial" ]]; then
        record_fail "SOA serial (${DOMAIN})" "could not parse serial"
    elif [[ "$serial" == "0" ]]; then
        record_fail "SOA serial (${DOMAIN})" "serial is 0"
    else
        record_pass "SOA serial (${DOMAIN})" "${serial}"
    fi
}

# ── 9. SOA Consistency ───────────────────────────────────────────────
test_soa_consistency() {
    if [[ -z "$ZONE_MASTER" ]]; then
        record_skip "SOA consistency" "ZONE_MASTER not set"
        return
    fi
    local serial_local serial_master
    # Get serial from configured server
    serial_local=$(dig_cmd +short "${DOMAIN}" SOA 2>/dev/null | awk '{print $3}') || true
    # Get serial from master
    serial_master=$("$DIG_CMD" "@${ZONE_MASTER}" +short "${DOMAIN}" SOA 2>/dev/null | awk '{print $3}') || true
    if [[ -z "$serial_local" || -z "$serial_master" ]]; then
        record_fail "SOA consistency" "could not retrieve serials (local=${serial_local:-?}, master=${serial_master:-?})"
        return
    fi
    if [[ "$serial_local" == "$serial_master" ]]; then
        record_pass "SOA consistency" "serial ${serial_local} matches across servers"
    else
        record_fail "SOA consistency" "serial mismatch — local=${serial_local}, master=${serial_master}"
    fi
}

# ── 10. Zone Transfer ────────────────────────────────────────────────
test_zone_transfer() {
    section_header "Zone Transfer"
    if [[ -z "$ZONE" ]]; then
        record_skip "Zone transfer (AXFR)" "ZONE not set"
        return
    fi
    local output exit_code=0
    output=$(dig_cmd AXFR "${ZONE}" 2>&1) || exit_code=$?
    # Check if transfer returned records
    local record_count
    record_count=$(echo "$output" | grep -c "^${ZONE}" 2>/dev/null) || record_count=0
    if [[ $record_count -gt 0 ]]; then
        record_pass "Zone transfer (${ZONE})" "${record_count} records transferred"
    elif echo "$output" | grep -qi "transfer failed\|refused\|REFUSED"; then
        record_pass "Zone transfer (${ZONE})" "AXFR refused (expected on production)"
    else
        record_fail "Zone transfer (${ZONE})" "transfer failed — ${output:0:100}"
    fi
}

# ── 11. DNSSEC Validation ────────────────────────────────────────────
test_dnssec_validation() {
    section_header "DNSSEC"
    if [[ -z "$DNSSEC_DOMAIN" ]]; then
        record_skip "DNSSEC validation" "DNSSEC_DOMAIN not set"
        return
    fi
    local output
    output=$(dig_cmd +dnssec +short "${DNSSEC_DOMAIN}" A 2>/dev/null) || true
    # Check for AD flag in full output
    local full_output
    full_output=$(dig_cmd +dnssec "${DNSSEC_DOMAIN}" A 2>/dev/null) || true
    if echo "$full_output" | grep -q "flags:.*ad"; then
        record_pass "DNSSEC validation (${DNSSEC_DOMAIN})" "AD flag set"
    elif [[ -n "$output" ]]; then
        record_fail "DNSSEC validation (${DNSSEC_DOMAIN})" "response received but AD flag not set"
    else
        record_fail "DNSSEC validation (${DNSSEC_DOMAIN})" "no response"
    fi
}

# ── 12. DNS-over-TLS ─────────────────────────────────────────────────
test_dot() {
    section_header "DNS-over-TLS"
    if [[ -z "$DOT_SERVER" ]]; then
        record_skip "DNS-over-TLS" "DOT_SERVER not set"
        return
    fi
    if ! has_cmd openssl; then
        record_skip "DNS-over-TLS" "openssl not installed"
        return
    fi
    local output exit_code=0
    output=$(echo "" | openssl s_client -connect "${DOT_SERVER}:853" -servername "${DOT_SERVER}" 2>&1) || exit_code=$?
    if echo "$output" | grep -qi "connected\|verify return"; then
        # Extract certificate info if available
        local cn
        cn=$(echo "$output" | grep -oP 'CN\s*=\s*\K[^,/]+' | head -1) || true
        record_pass "DNS-over-TLS (${DOT_SERVER}:853)" "TLS handshake OK${cn:+ — CN=${cn}}"
    else
        record_fail "DNS-over-TLS (${DOT_SERVER}:853)" "TLS handshake failed"
    fi
}

# ── 13. Custom Record Checks ─────────────────────────────────────────
test_custom_records() {
    if [[ -z "$TEST_RECORDS" ]]; then return; fi
    section_header "Custom Records"
    local IFS=','
    for entry in $TEST_RECORDS; do
        local name type expected
        name=$(echo "$entry" | cut -d: -f1)
        type=$(echo "$entry" | cut -d: -f2)
        expected=$(echo "$entry" | cut -d: -f3-)
        if [[ -z "$name" || -z "$type" ]]; then
            record_fail "Custom record" "invalid entry: ${entry}"
            continue
        fi
        local output
        output=$(dig_cmd +short "${name}" "${type}" 2>/dev/null) || true
        if [[ -z "$output" ]]; then
            record_fail "Custom record (${name} ${type})" "no record returned"
        elif [[ -n "$expected" ]]; then
            if echo "$output" | grep -q "$expected"; then
                record_pass "Custom record (${name} ${type})" "${output}"
            else
                record_fail "Custom record (${name} ${type})" "expected '${expected}', got '${output}'"
            fi
        else
            record_pass "Custom record (${name} ${type})" "${output}"
        fi
    done
}

# ── 14. Recursive Resolution ─────────────────────────────────────────
test_recursive_resolution() {
    section_header "Recursion"
    local output
    output=$(dig_cmd +short "google.com" A 2>/dev/null) || true
    if [[ -n "$output" ]]; then
        local first_ip
        first_ip=$(echo "$output" | head -1)
        record_pass "Recursive resolution (google.com)" "${first_ip}"
    else
        record_fail "Recursive resolution (google.com)" "could not resolve external domain"
    fi
}

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

usage() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Smoke-test DNS infrastructure. Checks resolution, zone transfers, SOA
consistency, DNSSEC, response time, and DNS-over-TLS.

Environment variables:
  DNS_SERVER        DNS server to query (default: system resolver)
  DOMAIN            Domain to test resolution (default: example.com)
  REVERSE_IP        IP for reverse (PTR) lookup (skip if empty)
  ZONE              Zone for AXFR transfer test (skip if empty)
  ZONE_MASTER       Second nameserver for SOA serial comparison (skip if empty)
  DNSSEC_DOMAIN     Domain to validate DNSSEC (skip if empty)
  DOT_SERVER        DNS-over-TLS server to test (skip if empty)
  MAX_RESPONSE_MS   Max acceptable response time in ms (default: 500)
  TEST_RECORDS      Comma-separated name:type:expected for custom checks
  OUTPUT_FORMAT     Output: text (default), tap
  COLOR             Color: auto (default), always, never
  VERBOSE           Show debug output (default: false)

Options:
  --server SERVER   Set DNS_SERVER
  --domain DOMAIN   Set DOMAIN
  --format FORMAT   Output: text (default), tap
  --verbose         Show debug output
  --no-color        Disable colored output
  --help            Show this help

Examples:
  DNS_SERVER=192.168.1.1 DOMAIN=corp.local ./$(basename "$0")
  DNSSEC_DOMAIN=example.com DOT_SERVER=1.1.1.1 ./$(basename "$0")
EOF
}

main() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --server)     DNS_SERVER="$2"; shift ;;
            --domain)     DOMAIN="$2"; shift ;;
            --format)     OUTPUT_FORMAT="$2"; shift ;;
            --verbose)    VERBOSE=true ;;
            --no-color)   COLOR=never ;;
            --help|-h)    usage; exit 0 ;;
            *)            err "Unknown option: $1"; usage; exit 1 ;;
        esac
        shift
    done

    setup_colors
    detect_dig

    START_TIME=$(date +%s)

    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        print_tap_header
    else
        print_header
    fi

    test_resolver_reachable
    test_forward_resolution
    test_aaaa_resolution
    test_mx_resolution
    test_reverse_lookup
    test_response_time
    test_authoritative_answer
    test_soa_serial
    test_soa_consistency
    test_zone_transfer
    test_dnssec_validation
    test_dot
    test_custom_records
    test_recursive_resolution

    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        print_tap_footer
    else
        print_summary
    fi

    [[ $FAIL -eq 0 ]] && exit 0 || exit 1
}
main "$@"