#!/bin/bash
# Manual WireGuard Setup Script
# Purpose: Step-by-step WireGuard installation and configuration
# This script shows what needs to be done - review before executing!

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' # No Color

print_step() {
    echo -e "${BLUE}[STEP]${NC} $1"
}

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

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

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

# ========================================
# Configuration
# ========================================

WG_INTERFACE="wg0"
WG_NETWORK="10.8.0.0/24"
WG_SERVER_IP="10.8.0.1"
WG_PORT="51820"
WG_CONFIG_DIR="/etc/wireguard"
WAN_INTERFACE="eth0"  # ANPASSEN an dein System!

# ========================================
# Pre-flight Checks
# ========================================

print_step "Pre-flight Checks"

if [ "$EUID" -ne 0 ]; then
    print_error "This script must be run as root"
    exit 1
fi

# Check if WireGuard is installed
if ! command -v wg &> /dev/null; then
    print_error "WireGuard is not installed"
    echo "Install with: apt update && apt install -y wireguard wireguard-tools qrencode nftables"
    exit 1
fi

print_success "Pre-flight checks passed"

# ========================================
# Step 1: Create WireGuard Directory
# ========================================

print_step "Creating WireGuard directory"

mkdir -p ${WG_CONFIG_DIR}
chmod 700 ${WG_CONFIG_DIR}

print_success "Directory created: ${WG_CONFIG_DIR}"

# ========================================
# Step 2: Generate Server Keys
# ========================================

print_step "Generating server keys"

cd ${WG_CONFIG_DIR}

if [ ! -f server_private.key ]; then
    wg genkey | tee server_private.key | wg pubkey > server_public.key
    chmod 600 server_private.key
    chmod 644 server_public.key
    print_success "Server keys generated"
else
    print_warning "Server keys already exist - skipping generation"
fi

SERVER_PRIVATE_KEY=$(cat server_private.key)
SERVER_PUBLIC_KEY=$(cat server_public.key)

echo ""
echo "Server Public Key: ${SERVER_PUBLIC_KEY}"
echo ""

# ========================================
# Step 3: Create WireGuard Configuration
# ========================================

print_step "Creating WireGuard configuration"

cat > ${WG_CONFIG_DIR}/${WG_INTERFACE}.conf <<EOF
[Interface]
# Server Configuration
PrivateKey = ${SERVER_PRIVATE_KEY}
Address = ${WG_SERVER_IP}/24
ListenPort = ${WG_PORT}

# Enable IP forwarding
PostUp = sysctl -w net.ipv4.ip_forward=1

# NAT Configuration with nftables
PostUp = nft add table inet wireguard
PostUp = nft add chain inet wireguard postrouting { type nat hook postrouting priority srcnat\; }
PostUp = nft add rule inet wireguard postrouting oifname "${WAN_INTERFACE}" ip saddr ${WG_NETWORK} masquerade

# Cleanup on shutdown
PostDown = nft delete table inet wireguard

# Peers will be added here via generate-client-config.sh
EOF

chmod 600 ${WG_CONFIG_DIR}/${WG_INTERFACE}.conf

print_success "Configuration created: ${WG_CONFIG_DIR}/${WG_INTERFACE}.conf"

# ========================================
# Step 4: Create nftables Firewall Rules
# ========================================

print_step "Creating nftables firewall rules"

cat > /etc/nftables.d/wireguard.nft <<'EOF'
#!/usr/sbin/nft -f

# WireGuard Host-based Firewall Configuration
# Purpose: Secure VPN access with admin service protection

table inet wireguard_firewall {
    # Define sets for efficient rule matching
    set vpn_network {
        type ipv4_addr
        flags interval
        elements = { 10.8.0.0/24 }
    }

    set admin_service_ports {
        type inet_service
        elements = {
            8080,  # Traefik Dashboard
            9090,  # Prometheus
            3001,  # Grafana
            9000,  # Portainer
            8001,  # Redis Insight
        }
    }

    set public_service_ports {
        type inet_service
        elements = {
            80,    # HTTP
            443,   # HTTPS
            22,    # SSH
        }
    }

    # Input chain - Control incoming connections
    chain input {
        type filter hook input priority filter; policy drop;

        # Allow established/related connections
        ct state established,related accept

        # Allow loopback
        iif lo accept

        # Allow ICMP (ping)
        ip protocol icmp accept
        ip6 nexthdr icmpv6 accept

        # Allow WireGuard port
        udp dport 51820 accept

        # Allow VPN network to access admin services
        ip saddr @vpn_network tcp dport @admin_service_ports accept

        # Allow public access to public services
        tcp dport @public_service_ports accept

        # Block public access to admin services (with logging)
        tcp dport @admin_service_ports counter log prefix "BLOCKED_ADMIN_SERVICE: " drop

        # Rate limit SSH to prevent brute force
        tcp dport 22 ct state new limit rate 10/minute accept

        # Drop everything else
        counter log prefix "BLOCKED_INPUT: " drop
    }

    # Forward chain - Control packet forwarding
    chain forward {
        type filter hook forward priority filter; policy drop;

        # Allow established/related connections
        ct state established,related accept

        # Allow VPN network to forward
        ip saddr @vpn_network accept

        # Drop everything else
        counter log prefix "BLOCKED_FORWARD: " drop
    }

    # Output chain - Allow all outgoing by default
    chain output {
        type filter hook output priority filter; policy accept;
    }
}
EOF

chmod 755 /etc/nftables.d/wireguard.nft

print_success "Firewall rules created: /etc/nftables.d/wireguard.nft"

# ========================================
# Step 5: Enable IP Forwarding
# ========================================

print_step "Enabling IP forwarding"

echo "net.ipv4.ip_forward=1" > /etc/sysctl.d/99-wireguard.conf
sysctl -p /etc/sysctl.d/99-wireguard.conf

print_success "IP forwarding enabled"

# ========================================
# Step 6: Apply nftables Rules
# ========================================

print_step "Applying nftables firewall rules"

if [ -f /etc/nftables.d/wireguard.nft ]; then
    nft -f /etc/nftables.d/wireguard.nft
    print_success "Firewall rules applied"
else
    print_error "Firewall rules file not found"
    exit 1
fi

# ========================================
# Step 7: Enable and Start WireGuard
# ========================================

print_step "Enabling and starting WireGuard service"

systemctl enable wg-quick@${WG_INTERFACE}
systemctl start wg-quick@${WG_INTERFACE}

print_success "WireGuard service enabled and started"

# ========================================
# Step 8: Verify Installation
# ========================================

print_step "Verifying installation"

echo ""
echo "WireGuard Status:"
wg show ${WG_INTERFACE}

echo ""
echo "Service Status:"
systemctl status wg-quick@${WG_INTERFACE} --no-pager

echo ""
echo "nftables Rules:"
nft list table inet wireguard_firewall

# ========================================
# Summary
# ========================================

echo ""
print_success "=========================================="
print_success "WireGuard Installation Complete!"
print_success "=========================================="
echo ""
echo "Server IP:        ${WG_SERVER_IP}"
echo "Listen Port:      ${WG_PORT}"
echo "VPN Network:      ${WG_NETWORK}"
echo "Interface:        ${WG_INTERFACE}"
echo ""
print_step "Next Steps:"
echo "  1. Generate client configs:"
echo "     cd /home/michael/dev/michaelschiemer/deployment/scripts"
echo "     sudo ./generate-client-config.sh <client-name>"
echo ""
echo "  2. Import client config on your device"
echo ""
echo "  3. Connect and test access to admin services:"
echo "     - Traefik Dashboard: https://10.8.0.1:8080"
echo "     - Prometheus: http://10.8.0.1:9090"
echo "     - Grafana: https://10.8.0.1:3001"
echo "     - Portainer: http://10.8.0.1:9000"
echo "     - Redis Insight: http://10.8.0.1:8001"
echo ""