#!/bin/bash

# Secure DNS Setup Script
# Configures systemd-resolved with DNSSEC + DoT and disables DHCP DNS
# Compatible with Ubuntu/Debian systems

set -euo pipefail

# Configuration URLs (customize these)
RESOLVED_CONF_URL="https://code.2pc.nexus/share/resolved.conf"
NETPLAN_CONF_URL="https://code.2pc.nexus/share/99-dns-override.yaml"

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging function
log() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

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

# Install required packages if missing
install_dependencies() {
    log "Checking and installing required dependencies..."
    
    # Update package list
    apt-get update -qq || {
        warn "Failed to update package list, continuing anyway"
    }
    
    # List of required packages
    local packages=("curl" "systemd-resolved" "netplan.io")
    local missing_packages=()
    
    # Check which packages are missing
    for package in "${packages[@]}"; do
        if ! dpkg -l | grep -q "^ii.*${package}"; then
            missing_packages+=("${package}")
        fi
    done
    
    # Install missing packages
    if [[ ${#missing_packages[@]} -gt 0 ]]; then
        log "Installing missing packages: ${missing_packages[*]}"
        apt-get install -y "${missing_packages[@]}" || {
            error "Failed to install required packages"
            exit 1
        }
    else
        success "All required packages are already installed"
    fi
}

# Check if systemd-resolved is available
check_systemd_resolved() {
    log "Checking systemd-resolved availability..."
    
    if ! systemctl list-unit-files | grep -q systemd-resolved; then
        error "systemd-resolved is not available on this system"
        exit 1
    fi
    
    # Enable and start systemd-resolved if not active
    if ! systemctl is-active --quiet systemd-resolved; then
        log "Enabling and starting systemd-resolved..."
        systemctl enable systemd-resolved
        systemctl start systemd-resolved
    fi
}

# Download configuration file with fallback
download_file() {
    local url="$1"
    local destination="$2"
    local description="$3"
    
    log "Downloading ${description}..."
    
    # Try curl first, then wget as fallback
    if command -v curl >/dev/null 2>&1; then
        curl -fsSL "${url}" -o "${destination}" || {
            error "Failed to download ${description} using curl"
            return 1
        }
    elif command -v wget >/dev/null 2>&1; then
        wget -q "${url}" -O "${destination}" || {
            error "Failed to download ${description} using wget"
            return 1
        }
    else
        error "Neither curl nor wget is available"
        return 1
    fi
    
    success "Downloaded ${description}"
}

# Create fallback configuration files
create_fallback_configs() {
    log "Creating fallback configuration files..."
    
    # Create systemd-resolved configuration
    cat > "/etc/systemd/resolved.conf.d/secure-dns.conf" << 'EOF'
[Resolve]
# DoT-enabled DNS servers with hostname for TLS certificate verification
DNS=9.9.9.9#dns.quad9.net 149.112.112.112#dns.quad9.net 1.1.1.1#cloudflare-dns.com 1.0.0.1#cloudflare-dns.com
FallbackDNS=8.8.8.8#dns.google 8.8.4.4#dns.google

# Security settings
DNSSEC=yes
DNSOverTLS=opportunistic

# Additional settings
Cache=yes
DNSStubListener=yes
EOF
    
    # Create netplan configuration
    cat > "/etc/netplan/99-dns-override.yaml" << 'EOF'
network:
  version: 2
  ethernets:
    eth0:
      dhcp4: true
      dhcp6: true
      dhcp4-overrides:
        use-dns: false
      dhcp6-overrides:
        use-dns: false
EOF
    
    success "Created fallback configuration files"
}

# Setup systemd-resolved configuration
setup_resolved_config() {
    log "Setting up systemd-resolved configuration..."
    
    # Create directory if it doesn't exist
    mkdir -p /etc/systemd/resolved.conf.d/
    
    # Backup existing configuration if it exists
    if [[ -f "/etc/systemd/resolved.conf.d/secure-dns.conf" ]]; then
        cp "/etc/systemd/resolved.conf.d/secure-dns.conf" \
           "/etc/systemd/resolved.conf.d/secure-dns.conf.backup.$(date +%Y%m%d_%H%M%S)"
        log "Backed up existing systemd-resolved configuration"
    fi
    
    # Try to download from URL, fallback to hardcoded config
    if ! download_file "${RESOLVED_CONF_URL}" "/etc/systemd/resolved.conf.d/secure-dns.conf" "systemd-resolved configuration"; then
        warn "Failed to download configuration, using fallback"
        create_fallback_configs
    fi
    
    # Set proper permissions
    chmod 644 /etc/systemd/resolved.conf.d/secure-dns.conf
    chown root:root /etc/systemd/resolved.conf.d/secure-dns.conf
}

# Setup netplan configuration
setup_netplan_config() {
    log "Setting up netplan configuration..."
    
    # Create netplan directory if it doesn't exist
    mkdir -p /etc/netplan/
    
    # Backup existing configuration if it exists
    if [[ -f "/etc/netplan/99-dns-override.yaml" ]]; then
        cp "/etc/netplan/99-dns-override.yaml" \
           "/etc/netplan/99-dns-override.yaml.backup.$(date +%Y%m%d_%H%M%S)"
        log "Backed up existing netplan configuration"
    fi
    
    # Try to download from URL, fallback to hardcoded config
    if ! download_file "${NETPLAN_CONF_URL}" "/etc/netplan/99-dns-override.yaml" "netplan configuration"; then
        warn "Failed to download netplan configuration, using fallback"
        # Fallback config is already created in create_fallback_configs
    fi
    
    # Set proper permissions (netplan requires 600)
    chmod 600 /etc/netplan/99-dns-override.yaml
    chown root:root /etc/netplan/99-dns-override.yaml
}

# Setup resolv.conf stub
setup_resolv_conf() {
    log "Setting up resolv.conf stub link..."
    
    # Backup existing resolv.conf
    if [[ -f "/etc/resolv.conf" ]] && [[ ! -L "/etc/resolv.conf" ]]; then
        cp "/etc/resolv.conf" "/etc/resolv.conf.backup.$(date +%Y%m%d_%H%M%S)"
        log "Backed up existing resolv.conf"
    fi
    
    # Create symlink to systemd-resolved stub
    ln -sf /run/systemd/resolve/stub-resolv.conf /etc/resolv.conf
    success "Configured resolv.conf stub link"
}

# Apply all configurations
apply_configurations() {
    log "Applying configurations..."
    
    # Apply netplan configuration
    if command -v netplan >/dev/null 2>&1; then
        netplan apply || {
            warn "netplan apply failed, continuing anyway"
        }
    else
        warn "netplan command not found, skipping netplan apply"
    fi
    
    # Restart systemd-resolved
    systemctl restart systemd-resolved || {
        error "Failed to restart systemd-resolved"
        exit 1
    }
    
    # Wait a moment for service to stabilize
    sleep 2
    
    # Flush DNS cache
    if command -v resolvectl >/dev/null 2>&1; then
        resolvectl flush-caches || {
            warn "Failed to flush DNS cache"
        }
    else
        warn "resolvectl command not found, skipping cache flush"
    fi
    
    success "Applied all configurations"
}

# Verify configuration
verify_configuration() {
    log "Verifying configuration..."
    
    # Check systemd-resolved status
    if ! systemctl is-active --quiet systemd-resolved; then
        error "systemd-resolved is not running"
        return 1
    fi
    
    # Check if resolvectl is available for status check
    if command -v resolvectl >/dev/null 2>&1; then
        log "Current DNS configuration:"
        resolvectl status || {
            warn "Failed to get resolvectl status"
        }
        
        # Test DNS resolution
        log "Testing DNS resolution..."
        if resolvectl query cloudflare.com >/dev/null 2>&1; then
            success "DNS resolution test passed"
        else
            warn "DNS resolution test failed"
        fi
    else
        warn "resolvectl not available for verification"
    fi
    
    # Check resolv.conf
    if [[ -L "/etc/resolv.conf" ]]; then
        local link_target
        link_target=$(readlink /etc/resolv.conf)
        if [[ "${link_target}" == "/run/systemd/resolve/stub-resolv.conf" ]]; then
            success "resolv.conf correctly linked to systemd-resolved stub"
        else
            warn "resolv.conf is linked to: ${link_target}"
        fi
    else
        warn "resolv.conf is not a symbolic link"
    fi
}

# Main execution flow
main() {
    log "Starting Secure DNS Setup..."
    log "This script will configure DNSSEC + DoT with systemd-resolved"
    
    # Pre-flight checks
    check_root
    install_dependencies
    check_systemd_resolved
    
    # Configuration setup
    setup_resolved_config
    setup_netplan_config
    setup_resolv_conf
    
    # Apply and verify
    apply_configurations
    verify_configuration
    
    success "Secure DNS setup completed successfully!"
    log "Your system now uses:"
    log "- DNSSEC validation (enforced)"
    log "- DNS over TLS (opportunistic)"
    log "- Secure DNS servers (Quad9, Cloudflare)"
    log "- DHCP DNS servers are ignored"
    
    log "You can verify the configuration with: resolvectl status"
}

# Handle script interruption
trap 'error "Script interrupted"; exit 1' INT TERM

# Run main function
main "$@"