#!/bin/bash
# WireGuard VPN Configuration Script for SAW
# Version: 1.0
# Date: 2026-04-02
#
# This script configures WireGuard VPN for the Secure Air-Gapped Workstation (SAW)
# It sets up the VPN tunnel, configures routing, and implements kill-switch functionality.
#
# Usage: sudo ./configure-vpn.sh
#
# SECURITY RATIONALE:
# WireGuard is used because it provides:
# 1. Modern, fast encryption (ChaCha20-Poly1305, Curve25519)
# 2. Simple configuration
# 3. Kernel-native implementation (faster than userspace)
# 4. Strong security properties
# 5. Easy to audit and verify
#
# The kill-switch ensures:
# - No traffic leaks when VPN is down
# - No direct internet access
# - All traffic must go through encrypted VPN tunnel
#
# Reference: https://www.wireguard.com/install/

set -e

# Configuration
VPN_CONFIG_FILE="/etc/wireguard/wg0.conf"
VPN_INTERFACE="wg0"
VPN_RULE_FILE="/etc/iproute2/rt_tables.d/51820-wg0"
KILL_SWITCH_SCRIPT="/usr/local/bin/wg-kill-switch.sh"
LOG_FILE="/var/log/wg-vpn-config.log"

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

# Print functions
log_info() {
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo -e "${GREEN}[$timestamp]${NC} ${BLUE}[INFO]${NC} $1" | tee -a "$LOG_FILE"
}

log_warning() {
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo -e "${YELLOW}[$timestamp]${NC} ${YELLOW}[WARNING]${NC} $1" | tee -a "$LOG_FILE"
}

log_error() {
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo -e "${RED}[$timestamp]${NC} ${RED}[ERROR]${NC} $1" | tee -a "$LOG_FILE"
}

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

# Check if WireGuard is installed
check_wireguard() {
    log_info "Checking WireGuard installation..."
    
    if ! command -v wg &> /dev/null; then
        log_warning "WireGuard not installed"
        log_info "Installing WireGuard..."
        dnf install -y wireguard-dkms wireguard-tools
    fi
    
    # Check if kernel module is loaded
    if ! lsmod | grep -q wireguard; then
        log_info "Loading WireGuard kernel module..."
        modprobe wireguard
    fi
    
    log_info "WireGuard is ready"
}

# Validate WireGuard configuration
validate_config() {
    log_info "Validating WireGuard configuration..."
    
    if [ ! -f "$VPN_CONFIG_FILE" ]; then
        log_error "WireGuard configuration not found: $VPN_CONFIG_FILE"
        log_info "Please create the configuration file first"
        
        # Show example configuration
        cat << 'EXAMPLE'
Example WireGuard Configuration:

[Interface]
# Your WireGuard private key (generated with: wg genkey)
PrivateKey = YOUR_PRIVATE_KEY_HERE
# VPN interface IP and subnet
Address = 10.0.0.2/24
# DNS servers (should be VPN gateway)
DNS = 10.0.0.1

[Peer]
# VPN server's public key
PublicKey = YOUR_SERVER_PUBLIC_KEY_HERE
# VPN server's endpoint (IP:port)
Endpoint = vpn.example.com:51820
# Which IPs to route through VPN (0.0.0.0/0 = all traffic)
AllowedIPs = 0.0.0.0/0, ::/0
EXAMPLE
        
        exit 1
    fi
    
    # Check private key permissions
    if [ -f "$VPN_CONFIG_FILE" ]; then
        local perm=$(stat -c %a "$VPN_CONFIG_FILE")
        if [ "$perm" != "600" ]; then
            log_warning "WireGuard config permissions: $perm (should be 600)"
            chmod 600 "$VPN_CONFIG_FILE"
            log_info "Fixed permissions to 600"
        fi
    fi
    
    # Validate configuration syntax
    if ! wg-quick parse-config "$VPN_CONFIG_FILE" >/dev/null 2>&1; then
        log_error "Invalid WireGuard configuration"
        wg-quick parse-config "$VPN_CONFIG_FILE"
        exit 1
    fi
    
    log_info "WireGuard configuration validated"
}

# Create routing table
create_routing_table() {
    log_info "Creating routing table for WireGuard..."
    
    # Create routing table definition
    cat > "$VPN_RULE_FILE" << 'ROUTES'
# WireGuard routing table
# Table 51820 is used for WireGuard traffic
51820	wg0
ROUTES
    
    # Add default route through WireGuard interface
    ip route add default dev "$VPN_INTERFACE" table 51820 2>/dev/null || true
    
    log_info "Routing table created"
}

# Configure iptables rules
configure_firewall() {
    log_info "Configuring firewall rules for WireGuard..."
    
    # Check if iptables is available
    if ! command -v iptables &> /dev/null; then
        log_warning "iptables not available, skipping"
        return
    fi
    
    # Allow WireGuard traffic
    iptables -A OUTPUT -o "$VPN_INTERFACE" -j ACCEPT 2>/dev/null || true
    iptables -A INPUT -i "$VPN_INTERFACE" -j ACCEPT 2>/dev/null || true
    
    # Block all other outbound traffic
    iptables -A OUTPUT ! -o lo ! -o "$VPN_INTERFACE" -j REJECT 2>/dev/null || true
    
    # Allow loopback
    iptables -A OUTPUT -o lo -j ACCEPT 2>/dev/null || true
    
    log_info "Firewall rules configured"
}

# Create kill-switch script
#
# SECURITY RATIONALE:
# The kill-switch is critical for security because it:
# 1. Prevents traffic leaks when VPN is down
# 2. Ensures all traffic goes through encrypted tunnel
# 3. Blocks DNS leaks (DNS only works through VPN)
# 4. Prevents data exfiltration when VPN is unavailable
# 5. Provides "fail closed" security model
#
# How it works:
# - Monitors VPN interface status
# - If VPN goes down, blocks all traffic (except loopback)
# - If VPN comes up, allows traffic through VPN
# - Can be configured to allow specific services (like update servers)
create_kill_switch() {
    log_info "Creating kill-switch script..."
    
    cat > "$KILL_SWITCH_SCRIPT" << 'KILLSCRIPT'
#!/bin/bash
# WireGuard Kill-Switch Script for SAW
# This script blocks all traffic when VPN is down

VPN_INTERFACE="wg0"
LOG_FILE="/var/log/wg-kill-switch.log"

# Function to enable kill-switch
enable_killswitch() {
    log_info "Enabling kill-switch..."
    
    # Flush existing rules
    iptables -F OUTPUT 2>/dev/null || true
    
    # Allow loopback
    iptables -A OUTPUT -o lo -j ACCEPT
    
    # Allow established connections (for cleanup)
    iptables -A OUTPUT -m state --state ESTABLISHED,RELATED -j ACCEPT
    
    # Block all other outbound traffic
    iptables -A OUTPUT -j REJECT
    
    log_info "Kill-switch enabled"
}

# Function to disable kill-switch
disable_killswitch() {
    log_info "Disabling kill-switch..."
    
    # Flush rules
    iptables -F OUTPUT 2>/dev/null || true
    
    # Allow all traffic (for VPN configuration)
    iptables -A OUTPUT -j ACCEPT
    
    log_info "Kill-switch disabled"
}

# Function to check VPN status
check_vpn() {
    if ip link show "$VPN_INTERFACE" >/dev/null 2>&1; then
        # VPN interface exists and is up
        if ip addr show "$VPN_INTERFACE" | grep -q "state UP"; then
            return 0  # VPN is up
        fi
    fi
    return 1  # VPN is down
}

# Function to log
log_info() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" >> "$LOG_FILE"
}

# Main logic
case "$1" in
    enable)
        enable_killswitch
        ;;
    disable)
        disable_killswitch
        ;;
    status)
        if check_vpn; then
            echo "VPN is UP"
            exit 0
        else
            echo "VPN is DOWN"
            exit 1
        fi
        ;;
    *)
        # Check VPN and act accordingly
        if check_vpn; then
            disable_killswitch
        else
            enable_killswitch
        fi
        ;;
esac
KILLSCRIPT
    
    chmod +x "$KILL_SWITCH_SCRIPT"
    log_info "Kill-switch script created"
}

# Configure systemd service for kill-switch
configure_killswitch_service() {
    log_info "Configuring kill-switch systemd service..."
    
    cat > /etc/systemd/system/wg-killswitch.service << 'KILLSERVICE'
[Unit]
Description=WireGuard Kill-Switch Service
After=network.target wg-quick@wg0.service
Wants=wg-quick@wg0.service

[Service]
Type=oneshot
ExecStart=/usr/local/bin/wg-kill-switch.sh
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
KILLSERVICE
    
    systemctl daemon-reload
    systemctl enable wg-killswitch
    log_info "Kill-switch service configured"
}

# Configure DNS to use VPN gateway
configure_dns() {
    log_info "Configuring DNS to use VPN gateway..."
    
    # Read VPN config to get DNS settings
    if [ -f "$VPN_CONFIG_FILE" ]; then
        local dns_servers=$(grep "^DNS =" "$VPN_CONFIG_FILE" | cut -d'=' -f2 | tr -d ' ')
        
        if [ -n "$dns_servers" ]; then
            # Configure resolv.conf
            echo "# WireGuard DNS configuration" > /etc/resolv.conf
            echo "# Generated by configure-vpn.sh" >> /etc/resolv.conf
            
            # Add DNS servers
            for server in $dns_servers; do
                echo "nameserver $server" >> /etc/resolv.conf
            done
            
            log_info "DNS configured to use VPN gateway"
        else
            log_warning "No DNS servers specified in WireGuard config"
        fi
    fi
    
    # Configure systemd-resolved to use VPN DNS
    if command -v systemctl &> /dev/null; then
        systemctl restart systemd-resolved 2>/dev/null || true
    fi
}

# Test VPN connectivity
test_vpn() {
    log_info "Testing VPN connectivity..."
    
    # Check if VPN interface is up
    if ip link show "$VPN_INTERFACE" | grep -q "state UP"; then
        log_info "VPN interface is UP"
        
        # Check connectivity
        if ping -c 3 -I "$VPN_INTERFACE" 8.8.8.8 >/dev/null 2>&1; then
            log_info "VPN connectivity test passed"
        else
            log_warning "VPN connectivity test failed (may need to wait for routing)"
        fi
    else
        log_warning "VPN interface is not UP"
        log_info "Run 'wg-quick up wg0' to start VPN"
    fi
}

# Display configuration
show_config() {
    echo ""
    echo "=== WireGuard Configuration ==="
    echo ""
    echo "Configuration file: $VPN_CONFIG_FILE"
    echo "VPN interface: $VPN_INTERFACE"
    echo ""
    
    if [ -f "$VPN_CONFIG_FILE" ]; then
        echo "Current configuration:"
        echo "----------------------------------------"
        grep -v "^#" "$VPN_CONFIG_FILE" | grep -v "^$" | sed 's/^/  /'
        echo "----------------------------------------"
    fi
    
    echo ""
    echo "Kill-switch script: $KILL_SWITCH_SCRIPT"
    echo ""
}

# Main execution
main() {
    echo ""
    echo "=========================================="
    echo "  WireGuard VPN Configuration"
    echo "  Fedora Kinoite SAW"
    echo "=========================================="
    echo ""
    
    check_root
    check_wireguard
    validate_config
    create_routing_table
    configure_firewall
    create_kill_switch
    configure_killswitch_service
    configure_dns
    show_config
    
    echo ""
    echo "=========================================="
    echo "  VPN Configuration Complete!"
    echo "=========================================="
    echo ""
    echo "Next steps:"
    echo "1. Edit WireGuard config: sudo nano $VPN_CONFIG_FILE"
    echo "2. Start VPN: sudo wg-quick up wg0"
    echo "3. Test VPN: ping -I wg0 8.8.8.8"
    echo "4. Check status: wg show"
    echo ""
    echo "To start VPN automatically on boot:"
    echo "  sudo systemctl enable wg-quick@wg0"
    echo ""
}

# Run main function
main "$@"