#!/usr/bin/env bash

#########################################################################################
#### dns-lookup.sh — Batch DNS lookups with record comparison across servers          ####
#### Query multiple record types and compare results across DNS servers               ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./dns-lookup.sh example.com google.com                                        ####
####   ./dns-lookup.sh --type MX --servers 8.8.8.8,1.1.1.1 example.com              ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
DNS_TIMEOUT="${DNS_TIMEOUT:-5}"
RECORD_TYPE="${RECORD_TYPE:-A}"
DNS_SERVERS=""
COMPARE="${COMPARE:-false}"
DOMAIN_FILE=""
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
DOMAINS=()
COUNT_TOTAL=0
COUNT_SUCCESS=0
COUNT_FAILED=0
COUNT_MISMATCH=0
DIG_CMD=""

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

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

# ── Helpers ───────────────────────────────────────────────────────────
section_header() {
    echo ""
    echo -e "  ${BOLD}${CYAN}── $1 ──${RESET}"
    echo ""
}

field() {
    printf "  ${BOLD}%-22s${RESET} %s\n" "$1" "$2"
}

field_color() {
    printf "  ${BOLD}%-22s${RESET} %b\n" "$1" "$2"
}

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

detect_dns_tool() {
    if command -v dig &>/dev/null; then
        DIG_CMD="dig"
    elif command -v nslookup &>/dev/null; then
        DIG_CMD="nslookup"
    else
        err "Neither dig nor nslookup found. Install dnsutils or bind-utils."
        exit 1
    fi
    verbose "Using DNS tool: ${DIG_CMD}"
}

query_dig() {
    local domain="$1"
    local rtype="$2"
    local server="${3:-}"

    local cmd_args=()
    if [[ -n "$server" ]]; then
        cmd_args+=("@${server}")
    fi
    cmd_args+=("$domain" "$rtype" "+short" "+time=${DNS_TIMEOUT}" "+tries=1")

    verbose "dig ${cmd_args[*]}"
    dig "${cmd_args[@]}" 2>/dev/null || echo ""
}

query_nslookup() {
    local domain="$1"
    local rtype="$2"
    local server="${3:-}"

    local result
    if [[ -n "$server" ]]; then
        result=$(nslookup -type="$rtype" -timeout="$DNS_TIMEOUT" "$domain" "$server" 2>/dev/null) || result=""
    else
        result=$(nslookup -type="$rtype" -timeout="$DNS_TIMEOUT" "$domain" 2>/dev/null) || result=""
    fi

    # Parse nslookup output — extract answer lines
    echo "$result" | awk '/^Name:|^Address:|answer:/{found=1} found && /^[^ \t]/' | grep -v "^Server:" | grep -v "^Name:" | awk '{print $NF}'
}

do_query() {
    local domain="$1"
    local rtype="$2"
    local server="${3:-}"

    if [[ "$DIG_CMD" == "dig" ]]; then
        query_dig "$domain" "$rtype" "$server"
    else
        query_nslookup "$domain" "$rtype" "$server"
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       LOOKUP LOGIC
# ══════════════════════════════════════════════════════════════════════

lookup_single() {
    local domain="$1"
    local rtype="$2"
    local server="${3:-system resolver}"
    local server_arg="${3:-}"

    COUNT_TOTAL=$((COUNT_TOTAL + 1))

    local result
    result=$(do_query "$domain" "$rtype" "$server_arg")

    if [[ -z "$result" ]]; then
        COUNT_FAILED=$((COUNT_FAILED + 1))
        printf "  %b%-30s %-6s %-18s %s%b\n" "$RED" "$domain" "$rtype" "$server" "NO RECORDS" "$RESET"
        return
    fi

    COUNT_SUCCESS=$((COUNT_SUCCESS + 1))

    # Get TTL if using dig
    local ttl="--"
    if [[ "$DIG_CMD" == "dig" && -n "$server_arg" ]]; then
        ttl=$(dig "@${server_arg}" "$domain" "$rtype" +noall +answer +time="${DNS_TIMEOUT}" +tries=1 2>/dev/null \
            | awk '{print $2}' | head -1 || echo "--")
    elif [[ "$DIG_CMD" == "dig" ]]; then
        ttl=$(dig "$domain" "$rtype" +noall +answer +time="${DNS_TIMEOUT}" +tries=1 2>/dev/null \
            | awk '{print $2}' | head -1 || echo "--")
    fi

    while IFS= read -r value; do
        [[ -z "$value" ]] && continue
        printf "  %-30s %-6s %-8s %-18s %s\n" "$domain" "$rtype" "$ttl" "$server" "$value"
        # Only print domain on first line
        domain=""
        ttl=""
    done <<< "$result"
}

lookup_compare() {
    local domain="$1"
    local rtype="$2"
    local -a servers_arr

    IFS=',' read -ra servers_arr <<< "$DNS_SERVERS"

    if [[ ${#servers_arr[@]} -lt 2 ]]; then
        warn "Compare mode requires at least 2 DNS servers (use --servers)"
        return
    fi

    local -a all_results=()
    local first_result=""

    for server in "${servers_arr[@]}"; do
        COUNT_TOTAL=$((COUNT_TOTAL + 1))

        local result
        result=$(do_query "$domain" "$rtype" "$server" | sort)

        if [[ -z "$result" ]]; then
            COUNT_FAILED=$((COUNT_FAILED + 1))
            printf "  %b%-30s %-6s %-18s %s%b\n" "$RED" "$domain" "$rtype" "$server" "NO RECORDS" "$RESET"
            all_results+=("FAILED")
            continue
        fi

        COUNT_SUCCESS=$((COUNT_SUCCESS + 1))
        all_results+=("$result")

        if [[ -z "$first_result" ]]; then
            first_result="$result"
        fi

        while IFS= read -r value; do
            [[ -z "$value" ]] && continue
            printf "  %-30s %-6s %-18s %s\n" "$domain" "$rtype" "$server" "$value"
            domain=""
        done <<< "$result"
    done

    # Check for mismatches
    local mismatch=false
    for r in "${all_results[@]}"; do
        if [[ "$r" != "$first_result" && "$r" != "FAILED" && "$first_result" != "FAILED" ]]; then
            mismatch=true
            break
        fi
    done

    if [[ "$mismatch" == "true" ]]; then
        COUNT_MISMATCH=$((COUNT_MISMATCH + 1))
        printf "  %b  ⚠ MISMATCH across servers for %s%b\n" "$RED" "$1" "$RESET"
    else
        printf "  %b  ✓ Consistent across servers for %s%b\n" "$GREEN" "$1" "$RESET"
    fi
    echo ""
}

# ══════════════════════════════════════════════════════════════════════
#                       INPUT PARSING
# ══════════════════════════════════════════════════════════════════════

parse_domain() {
    local entry="$1"
    entry=$(echo "$entry" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
    [[ -z "$entry" || "$entry" == \#* ]] && return
    DOMAINS+=("$entry")
}

load_domains_from_file() {
    local file="$1"
    if [[ ! -f "$file" ]]; then
        err "File not found: $file"
        exit 1
    fi
    while IFS= read -r line; do
        parse_domain "$line"
    done < "$file"
}

load_domains_from_stdin() {
    while IFS= read -r line; do
        parse_domain "$line"
    done
}

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

usage() {
    cat <<EOF
${SCRIPT_NAME} — Batch DNS lookups with record comparison

USAGE:
  ${SCRIPT_NAME} [OPTIONS] [DOMAIN ...]
  ${SCRIPT_NAME} --file domains.txt
  echo "example.com" | ${SCRIPT_NAME}

OPTIONS:
  --type TYPE           Record type: A, AAAA, MX, NS, TXT, CNAME, SOA, PTR
                        (default: A)
  --servers SERVERS     Comma-separated DNS servers to query
                        (default: system resolver)
  --compare             Compare results across multiple DNS servers
  --file FILE           Read domains from file (one per line)
  --verbose             Enable debug output
  --no-color            Disable colored output
  --help                Show this help

ENVIRONMENT VARIABLES:
  DNS_TIMEOUT           Query timeout in seconds (default: 5)

EXAMPLES:
  # Simple A record lookup
  ./dns-lookup.sh example.com

  # MX record lookup
  ./dns-lookup.sh --type MX example.com

  # Compare across DNS servers
  ./dns-lookup.sh --compare --servers 8.8.8.8,1.1.1.1,9.9.9.9 example.com

  # Batch lookup from file
  ./dns-lookup.sh --file domains.txt --type AAAA
EOF
}

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

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --type)
                RECORD_TYPE="${2^^}"; shift 2 ;;
            --servers)
                DNS_SERVERS="$2"; shift 2 ;;
            --compare)
                COMPARE="true"; shift ;;
            --file)
                DOMAIN_FILE="$2"; shift 2 ;;
            --verbose)
                VERBOSE="true"; shift ;;
            --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 ;;
            *)
                parse_domain "$1"; shift ;;
        esac
    done
}

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

main() {
    parse_args "$@"
    setup_colors
    detect_dns_tool

    # Load domains from file if specified
    if [[ -n "$DOMAIN_FILE" ]]; then
        load_domains_from_file "$DOMAIN_FILE"
    fi

    # Load from stdin if no domains yet and stdin is not a terminal
    if [[ ${#DOMAINS[@]} -eq 0 ]] && ! [[ -t 0 ]]; then
        load_domains_from_stdin
    fi

    if [[ ${#DOMAINS[@]} -eq 0 ]]; then
        err "No domains specified"
        echo "Run ${SCRIPT_NAME} --help for usage" >&2
        exit 1
    fi

    # Validate record type
    case "$RECORD_TYPE" in
        A|AAAA|MX|NS|TXT|CNAME|SOA|PTR) ;;
        *)
            err "Unsupported record type: ${RECORD_TYPE}"
            exit 1 ;;
    esac

    echo ""
    echo -e "${BOLD}DNS Lookup — ${RECORD_TYPE} Records${RESET}"
    echo -e "${DIM}$(date '+%Y-%m-%d %H:%M:%S %Z')${RESET}"
    echo -e "${DIM}Tool: ${DIG_CMD} | Timeout: ${DNS_TIMEOUT}s${RESET}"

    section_header "Results"

    if [[ "$COMPARE" == "true" ]]; then
        printf "  ${BOLD}%-30s %-6s %-18s %s${RESET}\n" "DOMAIN" "TYPE" "SERVER" "VALUE"
        printf "  %s\n" "$(printf '%.0s─' {1..85})"

        for domain in "${DOMAINS[@]}"; do
            lookup_compare "$domain" "$RECORD_TYPE"
        done
    else
        # Determine servers to query
        local -a servers_list
        if [[ -n "$DNS_SERVERS" ]]; then
            IFS=',' read -ra servers_list <<< "$DNS_SERVERS"
        else
            servers_list=("")
        fi

        printf "  ${BOLD}%-30s %-6s %-8s %-18s %s${RESET}\n" "DOMAIN" "TYPE" "TTL" "SERVER" "VALUE"
        printf "  %s\n" "$(printf '%.0s─' {1..90})"

        for domain in "${DOMAINS[@]}"; do
            for server in "${servers_list[@]}"; do
                lookup_single "$domain" "$RECORD_TYPE" "$server"
            done
        done
    fi

    section_header "Summary"
    field "Total lookups:" "$COUNT_TOTAL"
    field_color "Successful:" "${GREEN}${COUNT_SUCCESS}${RESET}"
    if [[ "$COUNT_FAILED" -gt 0 ]]; then
        field_color "Failed:" "${RED}${COUNT_FAILED}${RESET}"
    else
        field "Failed:" "$COUNT_FAILED"
    fi
    if [[ "$COMPARE" == "true" ]]; then
        if [[ "$COUNT_MISMATCH" -gt 0 ]]; then
            field_color "Mismatches:" "${RED}${COUNT_MISMATCH}${RESET}"
        else
            field_color "Mismatches:" "${GREEN}0${RESET}"
        fi
    fi

    echo ""
}

main "$@"