#!/bin/bash
set -e

# Setup script for RU VDS (Gateway)
# Run this script as root on the RU VDS server
#
# This script sets up:
# - WireGuard VPN (dual interface: users + DE tunnel)
# - dnsmasq for DNS resolution
# - nftables for firewall and routing
# - Policy routing for split-tunnel VPN

echo "========================================="
echo "RU VDS (Gateway) Setup"
echo "========================================="
echo ""

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    echo "ERROR: Please run as root"
    exit 1
fi

echo "[1/11] Updating system packages..."
apt update
apt upgrade -y

echo "[2/11] Installing required packages..."
apt install -y wireguard dnsmasq nftables qrencode curl bc

echo "[3/11] Disabling systemd-resolved (conflicts with dnsmasq)..."
systemctl disable --now systemd-resolved 2>/dev/null || true
rm -f /etc/resolv.conf
cat > /etc/resolv.conf << 'EOF'
nameserver 8.8.8.8
nameserver 1.1.1.1
EOF

echo "[4/11] Enabling IP forwarding..."
cat > /etc/sysctl.d/99-vpn.conf << 'EOF'
# Enable IP forwarding for VPN
net.ipv4.ip_forward = 1
EOF
sysctl -p /etc/sysctl.d/99-vpn.conf

echo "[5/11] Generating WireGuard keys..."
mkdir -p /etc/wireguard/keys
chmod 700 /etc/wireguard/keys

# Server key for user-facing interface
wg genkey | tee /etc/wireguard/keys/server.key | wg pubkey > /etc/wireguard/keys/server.pub

# Key for DE tunnel
wg genkey | tee /etc/wireguard/keys/de-tunnel.key | wg pubkey > /etc/wireguard/keys/de-tunnel.pub

chmod 600 /etc/wireguard/keys/*

echo "[6/11] Adding custom routing table..."
if ! grep -q "^200[[:space:]]*proxy" /etc/iproute2/rt_tables; then
    echo "200     proxy" >> /etc/iproute2/rt_tables
fi

echo "[7/11] Creating WireGuard configurations..."

# wg0 - user-facing
cat > /etc/wireguard/wg0.conf << 'EOF'
[Interface]
Address = 10.10.0.1/24
ListenPort = 51820
PrivateKey = __RU_SERVER_PRIVATE_KEY__
PostUp = /etc/wireguard/postup.sh
PostDown = /etc/wireguard/postdown.sh

# Client peers will be added below
# Use add-client.sh script to add new clients
EOF

# Replace private key placeholder
PRIVATE_KEY=$(cat /etc/wireguard/keys/server.key)
sed -i "s|__RU_SERVER_PRIVATE_KEY__|${PRIVATE_KEY}|g" /etc/wireguard/wg0.conf

# wg1 - DE tunnel
cat > /etc/wireguard/wg1.conf << 'EOF'
[Interface]
Address = 10.20.0.1/30
PrivateKey = __RU_DE_TUNNEL_PRIVATE_KEY__

[Peer]
# DE VDS (exit node)
PublicKey = __DE_SERVER_PUBLIC_KEY__
Endpoint = 194.31.173.178:51821
AllowedIPs = 10.10.0.0/24
PersistentKeepalive = 25
EOF

# Replace private key placeholder
DE_TUNNEL_KEY=$(cat /etc/wireguard/keys/de-tunnel.key)
sed -i "s|__RU_DE_TUNNEL_PRIVATE_KEY__|${DE_TUNNEL_KEY}|g" /etc/wireguard/wg1.conf

echo "[8/11] Creating WireGuard helper scripts..."

# PostUp script
cat > /etc/wireguard/postup.sh << 'EOF'
#!/bin/bash
set -e

#
# PostUp script for WireGuard wg0 interface
# Pure nftables solution - no iptables/ipset dependencies
#

# Load nftables rules (includes the 'direct' set and packet marking)
nft -f /etc/nftables.conf

# Add default route via DE tunnel for 'proxy' table
ip route add default via 10.20.0.2 dev wg1 table proxy 2>/dev/null || true

# Policy routing: packets with fwmark 0x1 use 'proxy' table
ip rule add from 10.10.0.0/24 fwmark 0x1 table proxy priority 100 2>/dev/null || true

echo "PostUp script completed successfully"
EOF

# PostDown script
cat > /etc/wireguard/postdown.sh << 'EOF'
#!/bin/bash

#
# PostDown script for WireGuard wg0 interface
# Pure nftables solution - no iptables/ipset dependencies
#

# Remove policy routing rule
ip rule del from 10.10.0.0/24 fwmark 0x1 table proxy priority 100 2>/dev/null || true

# Flush routing table
ip route flush table proxy 2>/dev/null || true

# Flush nftables vpn-routing table (keeps filter and nat rules intact)
nft flush table ip vpn-routing 2>/dev/null || true

echo "PostDown script completed"
EOF

chmod +x /etc/wireguard/postup.sh
chmod +x /etc/wireguard/postdown.sh

echo "[9/11] Creating nftables configuration..."
cat > /etc/nftables.conf << 'EOF'
#!/usr/sbin/nft -f
#
# RU VDS nftables configuration
#
# Routing approach:
# - Russian IP ranges loaded into 'direct' set by update-direct-routes.sh
# - nftables marks packets for policy routing
# - No iptables dependency
#

flush ruleset

table inet filter {
    chain input {
        type filter hook input priority 0; policy drop;

        # Allow established connections
        ct state established,related accept

        # Allow loopback
        iif lo accept

        # Allow SSH (adjust port if needed)
        tcp dport 22 accept

        # Allow WireGuard from anywhere (user connections)
        udp dport 51820 accept

        # Allow DNS from VPN clients only
        iifname "wg0" udp dport 53 accept
        iifname "wg0" tcp dport 53 accept

        # Allow ICMP
        icmp type echo-request accept
    }

    chain forward {
        type filter hook forward priority 0; policy drop;

        # Allow forwarding from user VPN
        iifname "wg0" accept

        # Allow forwarding from DE tunnel
        iifname "wg1" accept

        # Allow established connections
        ct state established,related accept
    }

    chain output {
        type filter hook output priority 0; policy accept;
    }
}

table ip vpn-routing {
    # Set for Russian IPs (direct routing, no proxy)
    # Populated by /etc/wireguard/update-direct-routes.sh
    # Auto-expires entries after 6 hours
    set direct {
        type ipv4_addr
        flags interval, timeout
        timeout 6h
    }

    # Packet marking chain for policy routing
    chain prerouting {
        type filter hook prerouting priority mangle; policy accept;

        # Only process traffic from VPN clients
        ip saddr != 10.10.0.0/24 return

        # Destinations in 'direct' set: no mark (direct routing)
        ip daddr @direct return

        # Everything else: mark for proxy routing via DE tunnel
        meta mark set 0x1
    }
}

table inet nat {
    chain postrouting {
        type nat hook postrouting priority 100;

        # NAT traffic going out to internet directly (not via wg1 tunnel)
        oifname != "wg0" oifname != "wg1" ip saddr 10.10.0.0/24 masquerade
    }
}
EOF

chmod +x /etc/nftables.conf

echo "[10/11] Configuring dnsmasq..."
cat > /etc/dnsmasq.d/vpn-routing.conf << 'EOF'
# dnsmasq configuration for VPN routing
#
# Note: Routing decisions are based on destination IP ranges,
# not DNS responses. Russian IP ranges are loaded into nftables
# by the update-direct-routes.sh script.

# Listen only on VPN interface
interface=wg0
bind-interfaces

# Upstream DNS servers
server=8.8.8.8
server=8.8.4.4
server=1.1.1.1

# Don't read /etc/resolv.conf
no-resolv

# Cache settings
cache-size=10000

# Log queries (optional, uncomment for debugging)
# log-queries
EOF

echo "[11/11] Creating Russian IP routes update script..."
cat > /etc/wireguard/update-direct-routes.sh << 'SCRIPT'
#!/bin/bash
#
# Downloads Russian IP ranges and adds them to the nftables 'direct' set
# These IPs will be routed directly instead of through the DE proxy
#
# Run this script:
# - Once during initial setup (after services are started)
# - Periodically via cron (weekly) to keep ranges updated
#
# Data source: RIPE NCC delegated statistics
#

set -e

RIPE_URL="https://ftp.ripe.net/pub/stats/ripencc/delegated-ripencc-extended-latest"
TEMP_FILE="/tmp/ripe-delegated.txt"
NFT_TABLE="ip vpn-routing"
NFT_SET="direct"

echo "Downloading RIPE delegation data..."
curl -s "$RIPE_URL" -o "$TEMP_FILE"

echo "Parsing Russian IP allocations..."

# Extract Russian IPv4 allocations and convert to CIDR notation
# Format: ripencc|RU|ipv4|start_ip|count|date|status
RU_RANGES=$(grep '|RU|ipv4|' "$TEMP_FILE" | while IFS='|' read -r registry cc type start count date status rest; do
    # Convert count to CIDR prefix length
    # count is number of IPs, prefix = 32 - log2(count)
    if [[ "$count" =~ ^[0-9]+$ ]] && [ "$count" -gt 0 ]; then
        prefix=$(echo "32 - l($count)/l(2)" | bc -l 2>/dev/null | cut -d. -f1)
        if [[ "$prefix" =~ ^[0-9]+$ ]] && [ "$prefix" -ge 8 ] && [ "$prefix" -le 32 ]; then
            echo "$start/$prefix"
        fi
    fi
done)

# Count ranges
RANGE_COUNT=$(echo "$RU_RANGES" | grep -c . || echo "0")
echo "Found $RANGE_COUNT Russian IP ranges"

echo "Flushing existing 'direct' set..."
nft flush set $NFT_TABLE $NFT_SET 2>/dev/null || true

echo "Adding ranges to nftables set (this may take a moment)..."
# Add in batches for efficiency
echo "$RU_RANGES" | while read -r cidr; do
    if [[ -n "$cidr" ]] && [[ "$cidr" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+/[0-9]+$ ]]; then
        nft add element $NFT_TABLE $NFT_SET { "$cidr" } 2>/dev/null || true
    fi
done

# Cleanup
rm -f "$TEMP_FILE"

echo ""
echo "Done. Russian IP ranges loaded into nftables."
echo "Traffic to these IPs will be routed directly (not through DE proxy)."
SCRIPT

chmod +x /etc/wireguard/update-direct-routes.sh

# Create clients directory
mkdir -p /etc/wireguard/clients

# Add cron job for weekly updates
echo "Setting up weekly cron job for IP range updates..."
cat > /etc/cron.weekly/update-vpn-routes << 'CRON'
#!/bin/bash
/etc/wireguard/update-direct-routes.sh >> /var/log/vpn-routes-update.log 2>&1
CRON
chmod +x /etc/cron.weekly/update-vpn-routes

echo ""
echo "========================================="
echo "Setup completed!"
echo "========================================="
echo ""
echo "IMPORTANT: Next steps"
echo ""
echo "1. Your RU VDS public keys are:"
echo ""
echo "   Server key (for clients):"
cat /etc/wireguard/keys/server.pub
echo ""
echo "   DE tunnel key (for DE VDS):"
cat /etc/wireguard/keys/de-tunnel.pub
echo ""
echo "2. You need to get the DE VDS public key"
echo ""
echo "3. Edit /etc/wireguard/wg1.conf and replace:"
echo "   __DE_SERVER_PUBLIC_KEY__ with the actual DE VDS public key"
echo ""
echo "4. Enable and start services:"
echo "   systemctl enable nftables dnsmasq wg-quick@wg0 wg-quick@wg1"
echo "   systemctl start dnsmasq"
echo "   systemctl start wg-quick@wg1"
echo "   systemctl start wg-quick@wg0"
echo ""
echo "5. Load Russian IP ranges (after tunnel is up):"
echo "   /etc/wireguard/update-direct-routes.sh"
echo ""
echo "6. Verify the tunnel:"
echo "   wg show"
echo "   ping 10.20.0.2"
echo ""
echo "7. Add clients using: /root/add-client.sh <client_name>"
echo ""