#!/bin/bash
set -e

WG_CONF="/etc/wireguard/wg0.conf"

# Parse the config file
PRIVATE_KEY=$(grep -oP 'PrivateKey\s*=\s*\K.*' "$WG_CONF" | tr -d ' ')
ADDRESS_V4=$(grep -oP 'Address\s*=\s*\K[^,]+' "$WG_CONF" | grep -v ':' | tr -d ' ')
ADDRESS_V6=$(grep -oP 'Address\s*=\s*\K.*' "$WG_CONF" | grep -oP '[^,]*::[^,]*' | tr -d ' ')
DNS_SERVERS=$(grep -oP 'DNS\s*=\s*\K.*' "$WG_CONF" | tr ',' '\n' | tr -d ' ')
PEER_PUBKEY=$(grep -oP 'PublicKey\s*=\s*\K.*' "$WG_CONF" | tr -d ' ')
ENDPOINT=$(grep -oP 'Endpoint\s*=\s*\K.*' "$WG_CONF" | tr -d ' ')
WG_ENDPOINT=$(echo "$ENDPOINT" | cut -d: -f1)
WG_PORT=$(echo "$ENDPOINT" | cut -d: -f2)

# Set DNS manually
if [ -n "$DNS_SERVERS" ]; then
    : > /etc/resolv.conf
    for dns in $DNS_SERVERS; do
        echo "nameserver $dns" >> /etc/resolv.conf
    done
fi

DEFAULT_IF=$(ip route | awk '/default/ {print $5; exit}')
DEFAULT_GW=$(ip route | awk '/default/ {print $3; exit}')

# --- IPv4 kill switch ---
iptables -A INPUT -i lo -j ACCEPT
iptables -A OUTPUT -o lo -j ACCEPT
iptables -A OUTPUT -d "$WG_ENDPOINT" -p udp --dport "$WG_PORT" -j ACCEPT
iptables -A INPUT -s "$WG_ENDPOINT" -p udp --sport "$WG_PORT" -j ACCEPT
iptables -A INPUT -i wg0 -j ACCEPT
iptables -A OUTPUT -o wg0 -j ACCEPT
iptables -A INPUT -i "$DEFAULT_IF" -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
# Allow container-to-container traffic
iptables -A INPUT -i eth0 -j ACCEPT
iptables -A OUTPUT -o eth0 -d 172.16.0.0/12 -j ACCEPT
iptables -A OUTPUT -o eth0 -d 10.0.0.0/8 -j ACCEPT
iptables -A OUTPUT -o eth0 -d 192.168.0.0/16 -j ACCEPT
iptables -A INPUT -j DROP
iptables -A OUTPUT -j DROP

# --- IPv6 kill switch ---
ip6tables -A INPUT -i lo -j ACCEPT
ip6tables -A OUTPUT -o lo -j ACCEPT
ip6tables -A INPUT -i wg0 -j ACCEPT
ip6tables -A OUTPUT -o wg0 -j ACCEPT
ip6tables -A INPUT -j DROP
ip6tables -A OUTPUT -j DROP

# --- Bring up WireGuard manually (no wg-quick) ---
ip link add wg0 type wireguard
echo "$PRIVATE_KEY" | wg set wg0 private-key /dev/stdin peer "$PEER_PUBKEY" endpoint "$ENDPOINT" allowed-ips 0.0.0.0/0,::/0

[ -n "$ADDRESS_V4" ] && ip addr add "$ADDRESS_V4" dev wg0
[ -n "$ADDRESS_V6" ] && ip addr add "$ADDRESS_V6" dev wg0

ip link set wg0 up

# Route the WireGuard endpoint via the real gateway (avoid routing loop)
ip route add "$WG_ENDPOINT"/32 via "$DEFAULT_GW" dev "$DEFAULT_IF"

# Route all other traffic through the tunnel
ip route add 0.0.0.0/1 dev wg0
ip route add 128.0.0.0/1 dev wg0

# IPv6 routes through the tunnel
if [ -n "$ADDRESS_V6" ]; then
    ip -6 route add ::/1 dev wg0
    ip -6 route add 8000::/1 dev wg0
fi

echo "VPN is up. Checking connection..."
curl -s --max-time 10 https://am.i.mullvad.net/connected || echo "Warning: could not verify Mullvad connection"
echo "Public IP: $(curl -s --max-time 10 https://am.i.mullvad.net/ip) | Location: $(curl -s --max-time 10 https://am.i.mullvad.net/country), $(curl -s --max-time 10 https://am.i.mullvad.net/city)"

echo "VPN gateway ready."

# Keep the container running, exit gracefully on SIGTERM
trap 'echo "Shutting down VPN..."; ip link del wg0 2>/dev/null; exit 0' SIGTERM SIGINT
sleep infinity &
wait $!