Configure WireGuard VPN with DNS filtering and ad blocking using Pi-hole and Unbound

#!/usr/bin/env bash

set -euo pipefail

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

# Global variables
SCRIPT_NAME=$(basename "$0")
UNBOUND_PORT=5335
PIHOLE_USER="pihole"

# Usage function
usage() {
    echo "Usage: $SCRIPT_NAME [options]"
    echo "Options:"
    echo "  -h, --help     Show this help message"
    echo "  --no-confirm   Skip confirmation prompts"
    echo ""
    echo "This script installs and configures WireGuard VPN with Pi-hole and Unbound"
    exit 1
}

# Logging functions
log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
log_warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }

# Cleanup function
cleanup() {
    local exit_code=$?
    if [ $exit_code -ne 0 ]; then
        log_error "Script failed. Cleaning up..."
        systemctl stop unbound 2>/dev/null || true
        systemctl disable unbound 2>/dev/null || true
        if command -v pihole >/dev/null 2>&1; then
            pihole uninstall --unattended 2>/dev/null || true
        fi
    fi
    exit $exit_code
}

trap cleanup ERR

# Parse arguments
NO_CONFIRM=false
while [[ $# -gt 0 ]]; do
    case $1 in
        -h|--help) usage ;;
        --no-confirm) NO_CONFIRM=true ;;
        *) log_error "Unknown option: $1"; usage ;;
    esac
    shift
done

# Check if running as root or with sudo
if [[ $EUID -ne 0 ]]; then
    log_error "This script must be run as root or with sudo"
    exit 1
fi

# Detect distribution
if [ -f /etc/os-release ]; then
    . /etc/os-release
    case "$ID" in
        ubuntu|debian) 
            PKG_MGR="apt"
            PKG_UPDATE="apt update && apt upgrade -y"
            PKG_INSTALL="apt install -y"
            IPTABLES_PKG="iptables-persistent"
            ;;
        almalinux|rocky|centos|rhel|ol|fedora) 
            PKG_MGR="dnf"
            PKG_UPDATE="dnf update -y"
            PKG_INSTALL="dnf install -y"
            IPTABLES_PKG="iptables-services"
            ;;
        amzn) 
            PKG_MGR="yum"
            PKG_UPDATE="yum update -y"
            PKG_INSTALL="yum install -y"
            IPTABLES_PKG="iptables-services"
            ;;
        *) 
            log_error "Unsupported distribution: $ID"
            exit 1
            ;;
    esac
else
    log_error "Cannot detect distribution. /etc/os-release not found."
    exit 1
fi

log_info "Detected distribution: $PRETTY_NAME"

# Confirmation prompt
if [ "$NO_CONFIRM" = false ]; then
    echo -e "${YELLOW}This script will install and configure:${NC}"
    echo "  • WireGuard VPN server"
    echo "  • Pi-hole DNS ad blocker"
    echo "  • Unbound recursive DNS resolver"
    echo ""
    read -p "Do you want to continue? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        log_info "Installation cancelled"
        exit 0
    fi
fi

# Step 1: Update system packages
echo -e "\n${BLUE}[1/6] Updating system packages...${NC}"
eval "$PKG_UPDATE"
log_success "System packages updated"

# Step 2: Install required packages
echo -e "\n${BLUE}[2/6] Installing required packages...${NC}"
if [ "$PKG_MGR" = "apt" ]; then
    $PKG_INSTALL wireguard unbound curl qrencode $IPTABLES_PKG
else
    $PKG_INSTALL wireguard-tools unbound curl qrencode $IPTABLES_PKG
    systemctl enable iptables
fi
log_success "Required packages installed"

# Step 3: Configure Unbound
echo -e "\n${BLUE}[3/6] Configuring Unbound DNS resolver...${NC}"

# Backup original config
if [ -f /etc/unbound/unbound.conf ]; then
    cp /etc/unbound/unbound.conf /etc/unbound/unbound.conf.backup
fi

# Create Unbound configuration
cat > /etc/unbound/unbound.conf << 'EOF'
server:
    verbosity: 0
    interface: 127.0.0.1
    port: 5335
    do-ip4: yes
    do-udp: yes
    do-tcp: yes
    do-ip6: no
    prefer-ip6: no
    harden-glue: yes
    harden-dnssec-stripped: yes
    use-caps-for-id: no
    edns-buffer-size: 1232
    prefetch: yes
    num-threads: 1
    so-rcvbuf: 1m
    private-address: 192.168.0.0/16
    private-address: 169.254.0.0/16
    private-address: 172.16.0.0/12
    private-address: 10.0.0.0/8
    private-address: fd00::/8
    private-address: fe80::/10
EOF

# Set proper permissions
chmod 644 /etc/unbound/unbound.conf
chown root:root /etc/unbound/unbound.conf

# Start and enable Unbound
systemctl enable unbound
systemctl restart unbound

# Verify Unbound is working
sleep 2
if ! systemctl is-active --quiet unbound; then
    log_error "Unbound failed to start"
    exit 1
fi

# Test DNS resolution
if ! dig @127.0.0.1 -p $UNBOUND_PORT example.com +short >/dev/null 2>&1; then
    log_warning "Unbound DNS test failed, but continuing..."
fi

log_success "Unbound configured and running on port $UNBOUND_PORT"

# Step 4: Install Pi-hole
echo -e "\n${BLUE}[4/6] Installing Pi-hole...${NC}"

# Create Pi-hole setup configuration for automated install
cat > /tmp/pihole-setup.conf << EOF
PIHOLE_INTERFACE=$(ip route | grep default | head -n1 | awk '{print $5}')
IPV4_ADDRESS=$(ip route get 1.1.1.1 | grep -oP 'src \K\S+')
IPV6_ADDRESS=
QUERY_LOGGING=true
INSTALL_WEB_SERVER=true
INSTALL_WEB_INTERFACE=true
LIGHTTPD_ENABLED=true
PIHOLE_DNS_1=127.0.0.1#5335
PIHOLE_DNS_2=
DNS_FQDN_REQUIRED=true
DNS_BOGUS_PRIV=true
DNSSEC=false
TEMPERATUREUNIT=C
WEBUIBOXEDLAYOUT=traditional
API_EXCLUDE_DOMAINS=
API_EXCLUDE_CLIENTS=
API_QUERY_LOG_SHOW=permittedonly
API_PRIVACY_MODE=false
EOF

# Install Pi-hole with automated configuration
curl -sSL https://install.pi-hole.net | bash /dev/stdin --unattended

log_success "Pi-hole installed"

# Step 5: Configure Pi-hole DNS settings
echo -e "\n${BLUE}[5/6] Configuring Pi-hole DNS settings...${NC}"

# Set Pi-hole to use Unbound
pihole -a setdns 127.0.0.1#5335

# Restart Pi-hole DNS service
pihole restartdns

# Verify Pi-hole is working
sleep 2
if ! systemctl is-active --quiet pihole-FTL; then
    log_error "Pi-hole failed to start"
    exit 1
fi

SERVER_IP=$(hostname -I | awk '{print $1}')
log_success "Pi-hole configured. Web interface available at: http://${SERVER_IP}/admin/"

# Step 6: Configure firewall and show WireGuard setup info
echo -e "\n${BLUE}[6/6] Finalizing configuration...${NC}"

# Configure firewall rules for different distributions
if command -v ufw >/dev/null 2>&1 && ufw status | grep -q "Status: active"; then
    ufw allow 51820/udp comment "WireGuard"
    ufw allow 53/udp comment "DNS"
    ufw allow 80/tcp comment "Pi-hole web"
elif command -v firewall-cmd >/dev/null 2>&1; then
    firewall-cmd --permanent --add-port=51820/udp --zone=public
    firewall-cmd --permanent --add-port=53/udp --zone=public
    firewall-cmd --permanent --add-port=80/tcp --zone=public
    firewall-cmd --reload
fi

# Generate WireGuard server keys if they don't exist
WG_DIR="/etc/wireguard"
mkdir -p "$WG_DIR"
chmod 755 "$WG_DIR"

if [ ! -f "$WG_DIR/server_private.key" ]; then
    wg genkey | tee "$WG_DIR/server_private.key" | wg pubkey > "$WG_DIR/server_public.key"
    chmod 600 "$WG_DIR/server_private.key"
    chmod 644 "$WG_DIR/server_public.key"
fi

log_success "Installation completed successfully!"

# Final verification
echo -e "\n${GREEN}=== Verification ===${NC}"
echo "✓ Unbound status: $(systemctl is-active unbound)"
echo "✓ Pi-hole status: $(systemctl is-active pihole-FTL)"
echo "✓ WireGuard tools: $(which wg)"

echo -e "\n${GREEN}=== Next Steps ===${NC}"
echo "1. Access Pi-hole web interface: http://${SERVER_IP}/admin/"
echo "2. Configure WireGuard server at: $WG_DIR/"
echo "3. Add blocklists in Pi-hole admin panel"
echo "4. Create WireGuard client configurations"

echo -e "\n${YELLOW}Note:${NC} Remember to configure your WireGuard server configuration"
echo "to use ${SERVER_IP}:53 as the DNS server for clients."