#!/bin/bash
set -e

# Setup script for RU VDS (Gateway)
# Run this script as root on the RU VDS server

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/10] Updating system packages..."
apt update
apt upgrade -y

echo "[2/10] Installing required packages..."
apt install -y wireguard dnsmasq nftables iptables ipset qrencode

echo "[3/10] 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/10] 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/10] 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/10] 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/10] 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/10] Creating WireGuard helper scripts..."

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

# Create ipsets for routing decisions
ipset create direct hash:net -exist
ipset flush direct

# 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

# Load nftables rules
nft -f /etc/nftables.conf

# Mark packets NOT going to 'direct' ipset with fwmark 0x1
iptables -t mangle -I PREROUTING -m set ! --match-set direct dst -s 10.10.0.0/24 -j MARK --set-mark 0x1

echo "PostUp script completed successfully"
EOF

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

# 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

# Remove iptables mangle rule
iptables -t mangle -F PREROUTING 2>/dev/null || true

# Destroy ipsets
ipset destroy direct 2>/dev/null || true

echo "PostDown script completed"
EOF

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

echo "[9/10] Creating nftables configuration..."
cat > /etc/nftables.conf << 'EOF'
#!/usr/sbin/nft -f

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 inet nat {
    chain postrouting {
        type nat hook postrouting priority 100;

        # NAT direct traffic going out main interface
        oifname != "wg0" oifname != "wg1" ip saddr 10.10.0.0/24 masquerade
    }
}
EOF

chmod +x /etc/nftables.conf

echo "[10/10] Configuring dnsmasq..."
cat > /etc/dnsmasq.d/vpn-routing.conf << 'EOF'
# 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

# Russian TLDs - add resolved IPs to 'direct' ipset
ipset=/ru/direct
ipset=/рф/direct
ipset=/su/direct

# All other domains will go through proxy (default routing)
EOF

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

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"
echo "   systemctl start dnsmasq"
echo "   systemctl start wg-quick@wg1"
echo "   systemctl start wg-quick@wg0"
echo ""
echo "5. Verify the tunnel:"
echo "   wg show"
echo "   ping 10.20.0.2"
echo ""
echo "6. Add clients using: /root/add-client.sh <client_name>"
echo ""