#!/bin/bash
################################################################################
# Script Name: install-nsd.sh
# Version: 1.1
# Description: Install and configure NSD (Name Server Daemon) authoritative
#              DNS server as primary or secondary nameserver with zone file
#              generation and nsd-control support
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
# Date: 2026-03-31
#
# Supported OS:
#   - Ubuntu / Debian
#   - RHEL / AlmaLinux / Rocky Linux / Fedora
#
# Usage:
#   sudo ./install-nsd.sh --zone example.com
#   sudo ./install-nsd.sh --zone example.com --zone example.org --with-control
#   sudo ./install-nsd.sh --secondary --master-ip 10.0.0.1 --zone example.com
#   sudo ./install-nsd.sh --uninstall
#   ./install-nsd.sh --dry-run --zone example.com
#
################################################################################

set -euo pipefail

# ============================================================================
# CONFIGURATION
# ============================================================================

readonly VERSION="1.0"
readonly SCRIPT_NAME="${0##*/}"
readonly LOG_FILE="/var/log/nsd-install.log"

# Defaults
ROLE="primary"
MASTER_IP=""
LISTEN_IP="0.0.0.0"
ZONES=()
WITH_CONTROL=false
UNINSTALL=false
DRY_RUN=false

# Paths (set after OS detection)
NSD_CONF=""
NSD_ZONE_DIR=""
NSD_USER=""
NSD_GROUP=""

# OS detection
OS=""
OS_FAMILY=""

# ============================================================================
# COLOR OUTPUT
# ============================================================================

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

print_success() { echo -e "${GREEN}  ✓ $1${NC}"; }
print_error()   { echo -e "${RED}  ✗ $1${NC}" >&2; }
print_warning() { echo -e "${YELLOW}  ⚠ $1${NC}"; }
print_info()    { echo -e "  → $1"; }

# ============================================================================
# LOGGING
# ============================================================================

log() {
    local msg="[$(date '+%Y-%m-%d %H:%M:%S')] $1"
    if [[ "$DRY_RUN" == "false" ]] && [[ -w "$(dirname "$LOG_FILE")" || -w "$LOG_FILE" ]]; then
        echo "$msg" >> "$LOG_FILE"
    fi
    echo "$msg"
}

log_error() {
    local msg="[$(date '+%Y-%m-%d %H:%M:%S')] ERROR: $1"
    if [[ "$DRY_RUN" == "false" ]] && [[ -w "$(dirname "$LOG_FILE")" || -w "$LOG_FILE" ]]; then
        echo "$msg" >> "$LOG_FILE"
    fi
    print_error "$1"
}

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

show_help() {
    cat <<EOF
Usage: $SCRIPT_NAME [OPTIONS]

Install and configure NSD authoritative DNS server.

ROLE OPTIONS:
    --primary               Configure as primary nameserver (default)
    --secondary             Configure as secondary nameserver
    --master-ip IP          Master IP for zone transfers (required with --secondary)

CONFIGURATION:
    --zone DOMAIN           Add a DNS zone (can be repeated)
    --ip IP                 IP address to listen on (default: 0.0.0.0)
    --with-control          Enable nsd-control and generate TLS keys

ACTIONS:
    --uninstall             Remove NSD, configuration, and zone files
    --dry-run               Show what would be done without executing
    -h, --help              Show this help message
    --version               Show version

EXAMPLES:
    $SCRIPT_NAME --zone example.com
    $SCRIPT_NAME --zone example.com --zone example.org --with-control
    $SCRIPT_NAME --secondary --master-ip 10.0.0.1 --zone example.com
    $SCRIPT_NAME --ip 192.168.1.10 --zone internal.local
    $SCRIPT_NAME --uninstall
    $SCRIPT_NAME --dry-run --zone example.com

EOF
    exit 0
}

show_version() {
    echo "$SCRIPT_NAME version $VERSION"
    exit 0
}

check_root() {
    if [[ $EUID -ne 0 ]]; then
        print_error "This script must be run as root (use sudo)"
        exit 1
    fi
}

# ============================================================================
# OS DETECTION
# ============================================================================

detect_os() {
    if [[ ! -f /etc/os-release ]]; then
        log_error "Cannot detect OS — /etc/os-release not found"
        exit 1
    fi

    OS=$({ grep ^ID= /etc/os-release || true; } | cut -d= -f2 | tr -d '"' | tr '[:upper:]' '[:lower:]')

    case "$OS" in
        ubuntu|debian)
            OS_FAMILY="debian"
            NSD_CONF="/etc/nsd/nsd.conf"
            NSD_ZONE_DIR="/etc/nsd/zones"
            NSD_USER="nsd"
            NSD_GROUP="nsd"
            ;;
        rhel|centos|rocky|almalinux|fedora)
            OS_FAMILY="rhel"
            NSD_CONF="/etc/nsd/nsd.conf"
            NSD_ZONE_DIR="/etc/nsd/zones"
            NSD_USER="nsd"
            NSD_GROUP="nsd"
            ;;
        *)
            log_error "Unsupported OS: $OS"
            echo "Supported: Ubuntu, Debian, RHEL, AlmaLinux, Rocky Linux, Fedora"
            exit 1
            ;;
    esac

    log "Detected OS: $OS (family: $OS_FAMILY)"
}

# ============================================================================
# SYSTEMD-RESOLVED HANDLING
# ============================================================================

disable_resolved() {
    if [[ "$OS" != "ubuntu" ]]; then
        return
    fi

    if systemctl is-active --quiet systemd-resolved 2>/dev/null; then
        print_warning "systemd-resolved is running and binds to port 53"
        if [[ "$DRY_RUN" == "true" ]]; then
            print_info "[DRY RUN] Would disable systemd-resolved"
            return
        fi

        log "Disabling systemd-resolved to free port 53"
        systemctl stop systemd-resolved
        systemctl disable systemd-resolved

        # Replace symlinked resolv.conf with a static one
        if [[ -L /etc/resolv.conf ]]; then
            rm -f /etc/resolv.conf
            echo "nameserver 1.1.1.1" > /etc/resolv.conf
            echo "nameserver 8.8.8.8" >> /etc/resolv.conf
        fi

        print_success "systemd-resolved disabled"
    fi
}

# ============================================================================
# INSTALLATION
# ============================================================================

install_nsd() {
    log "Installing NSD packages"

    if [[ "$DRY_RUN" == "true" ]]; then
        print_info "[DRY RUN] Would install NSD via package manager"
        return
    fi

    if [[ "$OS_FAMILY" == "debian" ]]; then
        apt-get update -qq
        apt-get install -y nsd
    elif [[ "$OS_FAMILY" == "rhel" ]]; then
        if [[ "$OS" != "fedora" ]]; then
            dnf install -y epel-release
        fi
        dnf install -y nsd
    fi

    if command -v nsd &>/dev/null; then
        print_success "NSD installed ($(nsd -v 2>&1 | head -1))"
    else
        log_error "NSD installation failed"
        exit 1
    fi
}

# ============================================================================
# CONFIGURATION
# ============================================================================

generate_nsd_conf() {
    log "Generating NSD configuration"

    if [[ "$DRY_RUN" == "true" ]]; then
        print_info "[DRY RUN] Would write $NSD_CONF"
        print_info "[DRY RUN] Role: $ROLE | Listen: $LISTEN_IP | Zones: ${ZONES[*]:-none}"
        return
    fi

    mkdir -p "$NSD_ZONE_DIR"

    # Build zone configuration blocks
    local zone_blocks=""
    for zone in "${ZONES[@]}"; do
        if [[ "$ROLE" == "primary" ]]; then
            zone_blocks+="
zone:
    name: \"${zone}\"
    zonefile: \"${NSD_ZONE_DIR}/${zone}.zone\"
    provide-xfr: 0.0.0.0/0 NOKEY
    notify: 0.0.0.0 NOKEY
"
        else
            zone_blocks+="
zone:
    name: \"${zone}\"
    zonefile: \"${NSD_ZONE_DIR}/${zone}.zone\"
    allow-notify: ${MASTER_IP} NOKEY
    request-xfr: AXFR ${MASTER_IP} NOKEY
"
        fi
    done

    # Build remote-control block
    local control_block=""
    if [[ "$WITH_CONTROL" == "true" ]]; then
        control_block="
remote-control:
    control-enable: yes
    control-interface: 127.0.0.1
    control-port: 8952
    server-key-file: \"/etc/nsd/nsd_server.key\"
    server-cert-file: \"/etc/nsd/nsd_server.pem\"
    control-key-file: \"/etc/nsd/nsd_control.key\"
    control-cert-file: \"/etc/nsd/nsd_control.pem\"
"
    fi

    cat > "$NSD_CONF" <<EOF
# NSD configuration
# Generated by $SCRIPT_NAME v$VERSION on $(date '+%Y-%m-%d %H:%M:%S')
# Role: $ROLE

server:
    do-ip4: yes
    do-ip6: no
    port: 53
    interface: ${LISTEN_IP}

    # Paths
    zonesdir: "/etc/nsd"
    database: ""
    logfile: "/var/log/nsd.log"
    pidfile: "/run/nsd/nsd.pid"

    # Performance
    server-count: 1
    tcp-count: 250
    tcp-timeout: 120

    # Security
    hide-version: yes
    refuse-any: yes
    verbosity: 1
${control_block}${zone_blocks}
EOF

    chown "${NSD_USER}:${NSD_GROUP}" "$NSD_CONF"
    chmod 640 "$NSD_CONF"
    print_success "NSD configuration written to $NSD_CONF"
}

# ============================================================================
# ZONE FILE GENERATION
# ============================================================================

generate_zone_files() {
    if [[ ${#ZONES[@]} -eq 0 ]]; then
        return
    fi

    if [[ "$ROLE" == "secondary" ]]; then
        log "Skipping zone file generation (secondary gets zones via transfer)"
        return
    fi

    log "Generating zone files"

    if [[ "$DRY_RUN" == "true" ]]; then
        for zone in "${ZONES[@]}"; do
            print_info "[DRY RUN] Would create zone file: ${NSD_ZONE_DIR}/${zone}.zone"
        done
        return
    fi

    mkdir -p "$NSD_ZONE_DIR"
    local serial
    serial=$(date +%Y%m%d01)

    for zone in "${ZONES[@]}"; do
        local zone_file="${NSD_ZONE_DIR}/${zone}.zone"

        if [[ -f "$zone_file" ]]; then
            print_warning "Zone file exists, skipping: $zone_file"
            continue
        fi

        cat > "$zone_file" <<EOF
; Zone file for ${zone}
; Generated by $SCRIPT_NAME on $(date '+%Y-%m-%d %H:%M:%S')
;
\$ORIGIN ${zone}.
\$TTL 3600

@   IN  SOA ns1.${zone}. admin.${zone}. (
        ${serial}   ; Serial (YYYYMMDDNN)
        3600        ; Refresh (1 hour)
        900         ; Retry (15 minutes)
        1209600     ; Expire (2 weeks)
        300         ; Minimum TTL (5 minutes)
    )

; Nameservers
@       IN  NS      ns1.${zone}.
@       IN  NS      ns2.${zone}.

; A records — update these with real IPs
ns1     IN  A       ${LISTEN_IP}
ns2     IN  A       ${LISTEN_IP}
@       IN  A       ${LISTEN_IP}
www     IN  CNAME   ${zone}.

; Mail (example — uncomment and edit)
; @     IN  MX  10  mail.${zone}.
; mail  IN  A       ${LISTEN_IP}
EOF

        chown "${NSD_USER}:${NSD_GROUP}" "$zone_file"
        chmod 640 "$zone_file"
        print_success "Zone file created: $zone_file"
    done
}

# ============================================================================
# NSD-CONTROL SETUP
# ============================================================================

setup_nsd_control() {
    if [[ "$WITH_CONTROL" == "false" ]]; then
        return
    fi

    log "Setting up nsd-control with TLS keys"

    if [[ "$DRY_RUN" == "true" ]]; then
        print_info "[DRY RUN] Would run nsd-control-setup to generate TLS keys"
        return
    fi

    if command -v nsd-control-setup &>/dev/null; then
        nsd-control-setup
        print_success "nsd-control TLS keys generated"
    else
        log_error "nsd-control-setup not found — cannot generate keys"
        exit 1
    fi
}

# ============================================================================
# PERMISSIONS & VALIDATION
# ============================================================================

set_permissions() {
    if [[ "$DRY_RUN" == "true" ]]; then
        print_info "[DRY RUN] Would set ownership to ${NSD_USER}:${NSD_GROUP}"
        return
    fi

    log "Setting file permissions"

    # Ensure NSD runtime directory exists
    mkdir -p /run/nsd
    chown "${NSD_USER}:${NSD_GROUP}" /run/nsd

    # Ensure zone directory ownership
    if [[ -d "$NSD_ZONE_DIR" ]]; then
        chown -R "${NSD_USER}:${NSD_GROUP}" "$NSD_ZONE_DIR"
    fi

    # Ensure log file exists and is writable
    touch /var/log/nsd.log
    chown "${NSD_USER}:${NSD_GROUP}" /var/log/nsd.log

    print_success "File permissions set"
}

validate_config() {
    if [[ "$DRY_RUN" == "true" ]]; then
        print_info "[DRY RUN] Would run nsd-checkconf $NSD_CONF"
        return
    fi

    log "Validating NSD configuration"

    if nsd-checkconf "$NSD_CONF"; then
        print_success "Configuration validated (nsd-checkconf passed)"
    else
        log_error "Configuration validation failed — check $NSD_CONF"
        exit 1
    fi
}

# ============================================================================
# SERVICE MANAGEMENT
# ============================================================================

start_nsd() {
    if [[ "$DRY_RUN" == "true" ]]; then
        print_info "[DRY RUN] Would enable and start NSD service"
        return
    fi

    log "Enabling and starting NSD service"

    systemctl enable nsd
    systemctl restart nsd

    sleep 2

    if systemctl is-active --quiet nsd; then
        print_success "NSD service is running"
    else
        log_error "NSD service failed to start"
        systemctl status nsd --no-pager >&2
        exit 1
    fi
}

# ============================================================================
# UNINSTALL
# ============================================================================

do_uninstall() {
    log "Uninstalling NSD"

    if [[ "$DRY_RUN" == "true" ]]; then
        print_info "[DRY RUN] Would stop and disable NSD service"
        print_info "[DRY RUN] Would remove NSD package and config files"
        return
    fi

    check_root

    # Stop service
    if systemctl is-active --quiet nsd 2>/dev/null; then
        systemctl stop nsd
        systemctl disable nsd
        print_success "NSD service stopped and disabled"
    fi

    # Remove package
    if [[ "$OS_FAMILY" == "debian" ]]; then
        apt-get remove --purge -y nsd
        apt-get autoremove -y
    elif [[ "$OS_FAMILY" == "rhel" ]]; then
        dnf remove -y nsd
    fi

    # Remove configuration and zone files
    rm -rf /etc/nsd
    rm -f /var/log/nsd.log

    print_success "NSD removed"
    log "Uninstall complete"
    exit 0
}

# ============================================================================
# ARGUMENT PARSING
# ============================================================================

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --primary)
                ROLE="primary"
                shift
                ;;
            --secondary)
                ROLE="secondary"
                shift
                ;;
            --master-ip)
                [[ -z "${2:-}" ]] && { log_error "--master-ip requires an IP address"; exit 1; }
                MASTER_IP="$2"
                shift 2
                ;;
            --zone)
                [[ -z "${2:-}" ]] && { log_error "--zone requires a domain name"; exit 1; }
                ZONES+=("$2")
                shift 2
                ;;
            --ip)
                [[ -z "${2:-}" ]] && { log_error "--ip requires an IP address"; exit 1; }
                LISTEN_IP="$2"
                shift 2
                ;;
            --with-control)
                WITH_CONTROL=true
                shift
                ;;
            --uninstall)
                UNINSTALL=true
                shift
                ;;
            --dry-run)
                DRY_RUN=true
                shift
                ;;
            -h|--help)
                show_help
                ;;
            --version)
                show_version
                ;;
            *)
                log_error "Unknown option: $1"
                echo "Use --help for usage information" >&2
                exit 1
                ;;
        esac
    done

    # Validate arguments
    if [[ "$ROLE" == "secondary" && -z "$MASTER_IP" ]]; then
        log_error "--secondary requires --master-ip"
        exit 1
    fi
}

# ============================================================================
# SUMMARY
# ============================================================================

print_summary() {
    echo ""
    echo "=== NSD Installation Complete ==="
    echo ""
    echo "  Role:            $ROLE"
    echo "  Listen address:  $LISTEN_IP:53"
    echo "  Config file:     $NSD_CONF"
    echo "  Zone directory:  $NSD_ZONE_DIR"
    [[ "$WITH_CONTROL" == "true" ]] && echo "  nsd-control:     enabled"
    echo ""

    if [[ ${#ZONES[@]} -gt 0 ]]; then
        echo "  Zones configured:"
        for zone in "${ZONES[@]}"; do
            echo "    • $zone"
        done
        echo ""
    fi

    echo "  Next steps:"
    if [[ "$ROLE" == "primary" ]]; then
        echo "    1. Edit zone files in $NSD_ZONE_DIR with real records"
        echo "    2. Update serial number after each change"
        echo "    3. Reload: nsd-control reload (or systemctl reload nsd)"
    else
        echo "    1. Ensure the master ($MASTER_IP) allows zone transfers"
        echo "    2. Check transfer status: nsd-control zonestatus"
    fi
    echo "    • Test: dig @${LISTEN_IP} <domain> A"
    echo "    • Logs: /var/log/nsd.log"
    echo ""
}

# ============================================================================
# MAIN
# ============================================================================

main() {
    parse_args "$@"

    echo ""
    echo "=== NSD Authoritative DNS Server Installer v${VERSION} ==="
    echo ""

    # Handle uninstall early
    if [[ "$UNINSTALL" == "true" ]]; then
        detect_os
        do_uninstall
    fi

    # Dry-run doesn't need root for preview
    if [[ "$DRY_RUN" == "false" ]]; then
        check_root
    fi

    detect_os
    disable_resolved
    install_nsd
    generate_nsd_conf
    generate_zone_files
    setup_nsd_control
    set_permissions
    validate_config
    start_nsd
    print_summary

    log "Installation complete"
}

main "$@"