#!/usr/bin/env bash
set -euo pipefail

# Wait for WireGuard to be up
for i in {1..30}; do
  ip link show wg0 >/dev/null 2>&1 && break
  sleep 1
done

# Tor transparent proxy ports (from torrc)
TOR_TRANS=9040
TOR_DNS=5353

# Allow WG UDP
iptables -C INPUT -p udp --dport 51820 -j ACCEPT || iptables -A INPUT -p udp --dport 51820 -j ACCEPT
# Loopback + established (idempotent-ish)
iptables -C INPUT -i lo -j ACCEPT || iptables -A INPUT -i lo -j ACCEPT
iptables -C OUTPUT -o lo -j ACCEPT || iptables -A OUTPUT -o lo -j ACCEPT
iptables -C INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT || iptables -A INPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
iptables -C OUTPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT || iptables -A OUTPUT -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT

# === Transparent redirection of wg0 traffic into Tor ===
# DNS to Tor DNSPort (UDP + TCP)
iptables -t nat -C PREROUTING -i wg0 -p udp --dport 53 -j REDIRECT --to-ports ${TOR_DNS} \
  || iptables -t nat -A PREROUTING -i wg0 -p udp --dport 53 -j REDIRECT --to-ports ${TOR_DNS}
iptables -t nat -C PREROUTING -i wg0 -p tcp --dport 53 -j REDIRECT --to-ports ${TOR_DNS} \
  || iptables -t nat -A PREROUTING -i wg0 -p tcp --dport 53 -j REDIRECT --to-ports ${TOR_DNS}

# All TCP to Tor TransPort
iptables -t nat -C PREROUTING -i wg0 -p tcp -j REDIRECT --to-ports ${TOR_TRANS} \
  || iptables -t nat -A PREROUTING -i wg0 -p tcp -j REDIRECT --to-ports ${TOR_TRANS}

# Fail closed: block anything from wg0 that didn't get redirected
iptables -C FORWARD -o wg0 -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT \
  || iptables -I FORWARD 1 -o wg0 -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
iptables -C FORWARD -i wg0 -p udp -j REJECT --reject-with icmp-port-unreachable \
  || iptables -I FORWARD 2 -i wg0 -p udp -j REJECT --reject-with icmp-port-unreachable
iptables -C FORWARD -i wg0 -p icmp -j REJECT --reject-with icmp-host-unreachable \
  || iptables -I FORWARD 3 -i wg0 -p icmp -j REJECT --reject-with icmp-host-unreachable
iptables -C FORWARD -i wg0 -j DROP || iptables -I FORWARD 4 -i wg0 -j DROP