#!/usr/bin/env bash
# ============================================================
# set-linux-dns.sh
# ============================================================
# Remotely executed via SSH (piped through stdin via plink/ssh).
# Detects Linux distro, changes DNS to new DC IP, tracks every
# modified file (base64-encoded before/after content), validates
# resolution, and outputs structured key=value lines for the
# PowerShell dispatcher (Invoke-LinuxDns.ps1) to parse.
#
# Usage (called by PowerShell dispatcher):
#   plink user@host -pw pass "sudo bash -s" < set-linux-dns.sh
#   plink user@host -pw pass "sudo bash -s -- --dry-run" < set-linux-dns.sh
#
# Arguments:
#   --dry-run     Read-only mode. No files modified, no services restarted.
#   --new-dns     Override new DNS IP (default: passed via env NEW_DNS)
#   --domain      Override validation domain (default: passed via env DNS_DOMAIN)
# ============================================================

set -euo pipefail

# ── Argument parsing ─────────────────────────────────────────────────────────
DRY_RUN=false
NEW_DNS="${NEW_DNS:-10.35.40.1}"
DNS_DOMAIN="${DNS_DOMAIN:-corp.local}"

while [[ $# -gt 0 ]]; do
    case "$1" in
        --dry-run)   DRY_RUN=true ;;
        --new-dns)   NEW_DNS="$2";    shift ;;
        --domain)    DNS_DOMAIN="$2"; shift ;;
        *) ;;
    esac
    shift
done

# ── Output helpers ───────────────────────────────────────────────────────────
# All output lines are key=value for easy PowerShell parsing.
# Multiline values are base64-encoded.

output_line() { printf '%s=%s\n' "$1" "$2"; }

b64_encode() {
    # base64 with no line wrapping — compatible with GNU and BSD base64
    if base64 --version 2>/dev/null | grep -q GNU; then
        printf '%s' "$1" | base64 -w 0
    else
        printf '%s' "$1" | base64 | tr -d '\n'
    fi
}

b64_file() {
    # base64-encode a file's content
    if base64 --version 2>/dev/null | grep -q GNU; then
        base64 -w 0 < "$1"
    else
        base64 < "$1" | tr -d '\n'
    fi
}

# ── Distro detection ─────────────────────────────────────────────────────────
if [[ ! -f /etc/os-release ]]; then
    output_line "STATUS"      "failed"
    output_line "ERROR"       "Cannot detect distro: /etc/os-release not found"
    exit 1
fi

source /etc/os-release
DISTRO_ID="${ID:-unknown}"
DISTRO_VERSION="${VERSION_ID:-unknown}"
DISTRO_PRETTY="${PRETTY_NAME:-$DISTRO_ID $DISTRO_VERSION}"

output_line "DISTRO"    "$DISTRO_PRETTY"
output_line "DISTRO_ID" "$DISTRO_ID"
output_line "DRY_RUN"   "$DRY_RUN"

# ── Capture current /etc/resolv.conf as reference ────────────────────────────
RESOLV_CONF_B64=""
if [[ -f /etc/resolv.conf ]]; then
    RESOLV_CONF_B64=$(b64_file /etc/resolv.conf)
fi

# ── Detect current DNS IPs from resolv.conf ──────────────────────────────────
OLD_DNS_LIST=$(grep -E '^nameserver' /etc/resolv.conf 2>/dev/null | awk '{print $2}' | tr '\n' ',' | sed 's/,$//')
if [[ -z "$OLD_DNS_LIST" ]]; then
    OLD_DNS_LIST="unknown"
fi
output_line "OLD_DNS" "$OLD_DNS_LIST"
output_line "NEW_DNS" "$NEW_DNS"

# ── File tracking state ──────────────────────────────────────────────────────
FILE_COUNT=0
declare -a MODIFIED_FILES=()
declare -a FILES_BEFORE=()
declare -a FILES_AFTER=()
SERVICE_RESTARTED=""

# Helper: capture file before modification
capture_before() {
    local filepath="$1"
    local idx="$2"
    MODIFIED_FILES[$idx]="$filepath"
    if [[ -f "$filepath" ]]; then
        FILES_BEFORE[$idx]=$(b64_file "$filepath")
    else
        FILES_BEFORE[$idx]=$(b64_encode "FILE_NOT_FOUND")
    fi
}

# Helper: capture file after modification
capture_after() {
    local filepath="$1"
    local idx="$2"
    if [[ -f "$filepath" ]]; then
        FILES_AFTER[$idx]=$(b64_file "$filepath")
    else
        FILES_AFTER[$idx]=$(b64_encode "FILE_NOT_FOUND_AFTER")
    fi
}

# ── Main DNS change logic — branched per distro ──────────────────────────────

# ────────────────────────────────────────────────────────────────────────────
# UBUNTU 22.04 / 24.04
# Primary: /etc/systemd/resolved.conf
# Secondary: /etc/netplan/*.yaml (if nameservers hardcoded)
# ────────────────────────────────────────────────────────────────────────────
change_dns_ubuntu() {
    local changed_resolved=false
    local changed_netplan=false

    # ── Layer 1: systemd-resolved.conf ───────────────────────────────────────
    local resolved_conf="/etc/systemd/resolved.conf"

    # Ensure the file exists — create a minimal one if missing
    if [[ ! -f "$resolved_conf" ]]; then
        if [[ "$DRY_RUN" == "false" ]]; then
            mkdir -p "$(dirname "$resolved_conf")"
            printf '[Resolve]\n' > "$resolved_conf"
        fi
    fi

    local idx=$FILE_COUNT
    capture_before "$resolved_conf" "$idx"

    if [[ "$DRY_RUN" == "false" ]]; then
        # Remove existing DNS= and FallbackDNS= lines, then set new values
        if grep -q '^\[Resolve\]' "$resolved_conf"; then
            # File has [Resolve] section — update in place
            sed -i '/^DNS=/d'          "$resolved_conf"
            sed -i '/^FallbackDNS=/d'  "$resolved_conf"
            sed -i '/^\[Resolve\]/a DNS='"$NEW_DNS" "$resolved_conf"
        else
            # Append [Resolve] section
            printf '\n[Resolve]\nDNS=%s\n' "$NEW_DNS" >> "$resolved_conf"
        fi
        capture_after "$resolved_conf" "$idx"
        changed_resolved=true
    else
        FILES_AFTER[$idx]=$(b64_encode "DRY_RUN_NO_CHANGE")
    fi

    FILE_COUNT=$((FILE_COUNT + 1))

    # ── Layer 2: Netplan YAML files ───────────────────────────────────────────
    # Check if any netplan file has a hardcoded nameservers block
    local netplan_dir="/etc/netplan"
    if [[ -d "$netplan_dir" ]]; then
        local netplan_files
        # Find yaml files that contain nameservers
        netplan_files=$(grep -rl "nameservers" "$netplan_dir" 2>/dev/null || true)

        for np_file in $netplan_files; do
            [[ -z "$np_file" ]] && continue

            local np_idx=$FILE_COUNT
            capture_before "$np_file" "$np_idx"

            if [[ "$DRY_RUN" == "false" ]]; then
                # Replace all IPs in the addresses list under nameservers block
                # This handles YAML indentation variants with python3 (available on Ubuntu)
                python3 - "$np_file" "$NEW_DNS" <<'PYEOF'
import sys, re

filepath = sys.argv[1]
new_dns  = sys.argv[2]

with open(filepath, 'r') as f:
    content = f.read()

# Replace nameservers addresses block
# Matches: addresses: [x.x.x.x, y.y.y.y] or multi-line list
# Single-line format: addresses: [x.x.x.x, y.y.y.y]
content = re.sub(
    r'(nameservers:\s*\n\s+addresses:\s*\[)[^\]]+(\])',
    r'\g<1>' + new_dns + r'\2',
    content
)
# Also handle: addresses: [x.x.x.x, y.y.y.y] on same line as nameservers
content = re.sub(
    r'(addresses:\s*\[)[^\]]+(\])',
    r'\g<1>' + new_dns + r'\2',
    content
)

with open(filepath, 'w') as f:
    f.write(content)

print("OK")
PYEOF
                capture_after "$np_file" "$np_idx"
                changed_netplan=true
            else
                FILES_AFTER[$np_idx]=$(b64_encode "DRY_RUN_NO_CHANGE")
            fi

            FILE_COUNT=$((FILE_COUNT + 1))
        done
    fi

    # ── Restart services ──────────────────────────────────────────────────────
    if [[ "$DRY_RUN" == "false" ]]; then
        if [[ "$changed_resolved" == "true" ]]; then
            systemctl restart systemd-resolved 2>&1 || true
            # Re-link resolv.conf if it's a symlink to stub resolver
            if [[ -L /etc/resolv.conf ]]; then
                ln -sf /run/systemd/resolve/resolv.conf /etc/resolv.conf 2>/dev/null || true
            fi
            SERVICE_RESTARTED="systemd-resolved"
        fi
        if [[ "$changed_netplan" == "true" ]]; then
            netplan apply 2>&1 || true
            SERVICE_RESTARTED="${SERVICE_RESTARTED:+$SERVICE_RESTARTED,}netplan"
        fi
    fi
}

# ────────────────────────────────────────────────────────────────────────────
# CENTOS 7 / ROCKY LINUX 7
# Target: /etc/sysconfig/network-scripts/ifcfg-<NIC with default route>
# Service: NetworkManager
# ────────────────────────────────────────────────────────────────────────────
change_dns_el7() {
    # Find NIC with default route
    local default_nic
    default_nic=$(ip route 2>/dev/null | grep '^default' | awk '{print $5}' | head -1)

    if [[ -z "$default_nic" ]]; then
        output_line "STATUS" "failed"
        output_line "ERROR"  "Cannot determine default route NIC"
        exit 1
    fi

    local ifcfg_file="/etc/sysconfig/network-scripts/ifcfg-${default_nic}"

    if [[ ! -f "$ifcfg_file" ]]; then
        output_line "STATUS" "failed"
        output_line "ERROR"  "ifcfg file not found: $ifcfg_file"
        exit 1
    fi

    output_line "NIC" "$default_nic"

    local idx=$FILE_COUNT
    capture_before "$ifcfg_file" "$idx"

    if [[ "$DRY_RUN" == "false" ]]; then
        # Remove existing DNS1, DNS2, DNS3 lines and insert new DNS1
        sed -i '/^DNS[0-9]*=/d' "$ifcfg_file"
        printf 'DNS1=%s\n' "$NEW_DNS" >> "$ifcfg_file"

        capture_after "$ifcfg_file" "$idx"
        FILE_COUNT=$((FILE_COUNT + 1))

        systemctl restart NetworkManager 2>&1 || true
        SERVICE_RESTARTED="NetworkManager"

        # Wait briefly for NM to re-apply settings
        sleep 2
    else
        FILES_AFTER[$idx]=$(b64_encode "DRY_RUN_NO_CHANGE")
        FILE_COUNT=$((FILE_COUNT + 1))
    fi
}

# ────────────────────────────────────────────────────────────────────────────
# ROCKY LINUX 8
# Detect NM connection file format: ifcfg (legacy) or keyfile (.nmconnection)
# ────────────────────────────────────────────────────────────────────────────
change_dns_rocky8() {
    # Get the active connection name and its backing file
    local nm_conn_name nm_conn_file nm_conn_format

    nm_conn_name=$(nmcli -t -f NAME,DEVICE con show --active 2>/dev/null | \
                   grep -v '^lo' | head -1 | cut -d: -f1)

    if [[ -z "$nm_conn_name" ]]; then
        # Fallback: try to detect via default route NIC
        local default_nic
        default_nic=$(ip route 2>/dev/null | grep '^default' | awk '{print $5}' | head -1)
        nm_conn_name=$(nmcli -t -f NAME,DEVICE con show --active 2>/dev/null | \
                       grep ":${default_nic}$" | cut -d: -f1 | head -1)
    fi

    if [[ -z "$nm_conn_name" ]]; then
        output_line "STATUS" "failed"
        output_line "ERROR"  "Cannot determine active NetworkManager connection"
        exit 1
    fi

    output_line "NM_CONNECTION" "$nm_conn_name"

    # Detect file format from nmcli
    nm_conn_file=$(nmcli -t -f NAME,FILENAME con show 2>/dev/null | \
                   grep "^${nm_conn_name}:" | cut -d: -f2-)

    if [[ "$nm_conn_file" == *"system-connections"* ]]; then
        nm_conn_format="keyfile"
    elif [[ "$nm_conn_file" == *"ifcfg"* ]]; then
        nm_conn_format="ifcfg"
    else
        # Default: check if ifcfg file exists for the default NIC
        local default_nic
        default_nic=$(ip route 2>/dev/null | grep '^default' | awk '{print $5}' | head -1)
        local ifcfg_candidate="/etc/sysconfig/network-scripts/ifcfg-${default_nic}"
        if [[ -f "$ifcfg_candidate" ]]; then
            nm_conn_file="$ifcfg_candidate"
            nm_conn_format="ifcfg"
        else
            nm_conn_format="keyfile"
            nm_conn_file="/etc/NetworkManager/system-connections/${nm_conn_name}.nmconnection"
        fi
    fi

    output_line "NM_FORMAT" "$nm_conn_format"
    output_line "NM_FILE"   "$nm_conn_file"

    local idx=$FILE_COUNT
    capture_before "$nm_conn_file" "$idx"

    if [[ "$DRY_RUN" == "false" ]]; then
        if [[ "$nm_conn_format" == "ifcfg" ]]; then
            # Legacy ifcfg format — same as EL7
            sed -i '/^DNS[0-9]*=/d' "$nm_conn_file"
            printf 'DNS1=%s\n' "$NEW_DNS" >> "$nm_conn_file"
            capture_after "$nm_conn_file" "$idx"
            FILE_COUNT=$((FILE_COUNT + 1))

            systemctl restart NetworkManager 2>&1 || true
            SERVICE_RESTARTED="NetworkManager"
            sleep 2
        else
            # Keyfile format — use nmcli to modify then reload
            # nmcli modifies the .nmconnection file directly
            capture_before "$nm_conn_file" "$idx"

            nmcli con mod "$nm_conn_name" ipv4.dns "$NEW_DNS" 2>&1
            nmcli con mod "$nm_conn_name" ipv4.ignore-auto-dns "yes" 2>&1 || true

            # Bring the connection up to apply changes
            nmcli con up "$nm_conn_name" 2>&1 || true

            capture_after "$nm_conn_file" "$idx"
            FILE_COUNT=$((FILE_COUNT + 1))

            SERVICE_RESTARTED="NetworkManager(nmcli)"
            sleep 2
        fi
    else
        FILES_AFTER[$idx]=$(b64_encode "DRY_RUN_NO_CHANGE")
        FILE_COUNT=$((FILE_COUNT + 1))
    fi
}

# ── Dispatch to distro-specific function ─────────────────────────────────────
case "$DISTRO_ID" in
    ubuntu)
        change_dns_ubuntu
        ;;
    centos)
        change_dns_el7
        ;;
    rocky)
        # Rocky 8+ uses different NM handling
        if [[ "${DISTRO_VERSION%%.*}" -ge 8 ]]; then
            change_dns_rocky8
        else
            change_dns_el7
        fi
        ;;
    rhel|almalinux)
        # Treat like Rocky — same NM stack
        if [[ "${DISTRO_VERSION%%.*}" -ge 8 ]]; then
            change_dns_rocky8
        else
            change_dns_el7
        fi
        ;;
    *)
        output_line "STATUS" "failed"
        output_line "ERROR"  "Unsupported distro: $DISTRO_ID $DISTRO_VERSION"
        exit 1
        ;;
esac

# ── Output file tracking results ─────────────────────────────────────────────
output_line "FILE_COUNT"       "$FILE_COUNT"
output_line "SERVICE_RESTARTED" "$SERVICE_RESTARTED"

for ((i=0; i<FILE_COUNT; i++)); do
    n=$((i+1))
    output_line "MODIFIED_FILE_${n}"  "${MODIFIED_FILES[$i]}"
    output_line "FILE_BEFORE_${n}"    "${FILES_BEFORE[$i]}"
    output_line "FILE_AFTER_${n}"     "${FILES_AFTER[$i]}"
done

# ── Capture resolv.conf after changes ────────────────────────────────────────
RESOLV_CONF_AFTER_B64=""
if [[ -f /etc/resolv.conf ]]; then
    # Brief wait to let NM/resolved update resolv.conf
    if [[ "$DRY_RUN" == "false" ]]; then
        sleep 1
    fi
    RESOLV_CONF_AFTER_B64=$(b64_file /etc/resolv.conf)
fi

output_line "RESOLV_CONF_BEFORE" "$RESOLV_CONF_B64"
output_line "RESOLV_CONF_AFTER"  "$RESOLV_CONF_AFTER_B64"

# ── DNS validation (Task 7) ───────────────────────────────────────────────────
# Validate that the new DNS server can resolve the domain.
# Uses nslookup (universally available) with fallback to dig/host.
VALIDATION_STATUS="unknown"
VALIDATION_OUTPUT=""

do_validation() {
    local dns_server="$1"
    local domain="$2"

    if command -v nslookup &>/dev/null; then
        local nslookup_out
        nslookup_out=$(nslookup "$domain" "$dns_server" 2>&1) || true
        if echo "$nslookup_out" | grep -qiE 'Address.*[0-9]+\.[0-9]+|answer:'; then
            VALIDATION_STATUS="success"
        else
            VALIDATION_STATUS="failed"
        fi
        VALIDATION_OUTPUT=$(b64_encode "$nslookup_out")

    elif command -v dig &>/dev/null; then
        local dig_out
        dig_out=$(dig "@${dns_server}" "$domain" A +short 2>&1) || true
        if [[ -n "$dig_out" ]] && echo "$dig_out" | grep -qE '^[0-9]+\.[0-9]+'; then
            VALIDATION_STATUS="success"
        else
            VALIDATION_STATUS="failed"
        fi
        VALIDATION_OUTPUT=$(b64_encode "$dig_out")

    elif command -v host &>/dev/null; then
        local host_out
        host_out=$(host "$domain" "$dns_server" 2>&1) || true
        if echo "$host_out" | grep -qi 'has address'; then
            VALIDATION_STATUS="success"
        else
            VALIDATION_STATUS="failed"
        fi
        VALIDATION_OUTPUT=$(b64_encode "$host_out")

    else
        VALIDATION_STATUS="skipped"
        VALIDATION_OUTPUT=$(b64_encode "No DNS lookup tool found (nslookup/dig/host)")
    fi
}

if [[ "$DRY_RUN" == "false" ]]; then
    do_validation "$NEW_DNS" "$DNS_DOMAIN"
else
    # In dry-run: validate against current DNS to show current resolution state
    CURRENT_DNS=$(grep -E '^nameserver' /etc/resolv.conf 2>/dev/null | awk '{print $2}' | head -1)
    if [[ -n "$CURRENT_DNS" ]]; then
        do_validation "$CURRENT_DNS" "$DNS_DOMAIN"
        VALIDATION_STATUS="dryrun_${VALIDATION_STATUS}"
    else
        VALIDATION_STATUS="dryrun_skipped"
        VALIDATION_OUTPUT=$(b64_encode "No nameserver in resolv.conf to test against")
    fi
fi

output_line "VALIDATION"        "$VALIDATION_STATUS"
output_line "VALIDATION_OUTPUT" "$VALIDATION_OUTPUT"

# ── Final status ──────────────────────────────────────────────────────────────
if [[ "$DRY_RUN" == "true" ]]; then
    output_line "STATUS" "dryrun"
elif [[ "$FILE_COUNT" -eq 0 ]]; then
    output_line "STATUS"  "failed"
    output_line "ERROR"   "No files were modified — distro logic may not have run"
else
    output_line "STATUS" "success"
fi

exit 0