#!/usr/bin/env bash
# setup-wireguard-routing.sh
# Idempotent WireGuard split-tunnel routing helper.
# Detects iptables/nftables and optional UFW/Firewalld to configure forwarding + NAT.

set -euo pipefail

WAN_IF=${WAN_IF:-eth0}
WG_IF=${WG_IF:-wg0}
WG_NET=${WG_NET:-10.8.0.0/24}
WG_ADDR=${WG_ADDR:-10.8.0.1/24}
WG_PORT=${WG_PORT:-51820}
EXTRA_NETS_DEFAULT="192.168.178.0/24 172.20.0.0/16"
EXTRA_NETS="${EXTRA_NETS:-$EXTRA_NETS_DEFAULT}"
FIREWALL_BACKEND=${FIREWALL_BACKEND:-auto}
FIREWALLD_ZONE=${FIREWALLD_ZONE:-public}

read -r -a EXTRA_NETS_ARRAY <<< "${EXTRA_NETS}"

abort() {
  echo "Error: $1" >&2
  exit 1
}

require_root() {
  if [[ "${EUID}" -ne 0 ]]; then
    abort "please run as root (sudo ./setup-wireguard-routing.sh)"
  fi
}

detect_backend() {
  case "${FIREWALL_BACKEND}" in
    iptables|nftables) echo "${FIREWALL_BACKEND}"; return 0 ;;
    auto)
      if command -v nft >/dev/null 2>&1; then
        echo "nftables"; return 0
      fi
      if command -v iptables >/dev/null 2>&1; then
        echo "iptables"; return 0
      fi
      ;;
  esac
  abort "no supported firewall backend found (install iptables or nftables)"
}

ensure_sysctl() {
  local sysctl_file="/etc/sysctl.d/99-${WG_IF}-forward.conf"
  cat <<EOF > "${sysctl_file}"
# Managed by setup-wireguard-routing.sh
net.ipv4.ip_forward=1
EOF
  sysctl -w net.ipv4.ip_forward=1 >/dev/null
  sysctl --system >/dev/null
}

apply_iptables() {
  iptables -t nat -C POSTROUTING -s "${WG_NET}" -o "${WAN_IF}" -j MASQUERADE 2>/dev/null || \
    iptables -t nat -A POSTROUTING -s "${WG_NET}" -o "${WAN_IF}" -j MASQUERADE

  iptables -C FORWARD -i "${WG_IF}" -s "${WG_NET}" -o "${WAN_IF}" -j ACCEPT 2>/dev/null || \
    iptables -A FORWARD -i "${WG_IF}" -s "${WG_NET}" -o "${WAN_IF}" -j ACCEPT

  iptables -C FORWARD -o "${WG_IF}" -d "${WG_NET}" -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT 2>/dev/null || \
    iptables -A FORWARD -o "${WG_IF}" -d "${WG_NET}" -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT

  for net in "${EXTRA_NETS_ARRAY[@]}"; do
    [[ -z "${net}" ]] && continue
    iptables -C FORWARD -i "${WG_IF}" -d "${net}" -j ACCEPT 2>/dev/null || \
      iptables -A FORWARD -i "${WG_IF}" -d "${net}" -j ACCEPT
    iptables -C FORWARD -s "${net}" -o "${WG_IF}" -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT 2>/dev/null || \
      iptables -A FORWARD -s "${net}" -o "${WG_IF}" -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT
  done
}

apply_nftables() {
  local table="inet wireguard_${WG_IF}"
  nft list table ${table} >/dev/null 2>&1 || nft add table ${table}

  nft list chain ${table} postrouting >/dev/null 2>&1 || \
    nft add chain ${table} postrouting '{ type nat hook postrouting priority srcnat; }'

  nft list chain ${table} forward >/dev/null 2>&1 || \
    nft add chain ${table} forward '{ type filter hook forward priority filter; policy accept; }'

  nft list chain ${table} postrouting | grep -q "oifname \"${WAN_IF}\" ip saddr ${WG_NET}" || \
    nft add rule ${table} postrouting oifname "${WAN_IF}" ip saddr ${WG_NET} masquerade

  nft list chain ${table} forward | grep -q "iifname \"${WG_IF}\" ip saddr ${WG_NET}" || \
    nft add rule ${table} forward iifname "${WG_IF}" ip saddr ${WG_NET} counter accept

  nft list chain ${table} forward | grep -q "oifname \"${WG_IF}\" ip daddr ${WG_NET}" || \
    nft add rule ${table} forward oifname "${WG_IF}" ip daddr ${WG_NET} ct state established,related counter accept

  for net in "${EXTRA_NETS_ARRAY[@]}"; do
    [[ -z "${net}" ]] && continue
    nft list chain ${table} forward | grep -q "iifname \"${WG_IF}\" ip daddr ${net}" || \
      nft add rule ${table} forward iifname "${WG_IF}" ip daddr ${net} counter accept
  done
}

configure_ufw() {
  if command -v ufw >/dev/null 2>&1 && ufw status | grep -iq "Status: active"; then
    sed -i 's/^DEFAULT_FORWARD_POLICY=.*/DEFAULT_FORWARD_POLICY="ACCEPT"/' /etc/default/ufw
    ufw allow "${WG_PORT}"/udp >/dev/null
    ufw route allow in on "${WG_IF}" out on "${WAN_IF}" to any >/dev/null 2>&1 || true
    for net in "${EXTRA_NETS_ARRAY[@]}"; do
      [[ -z "${net}" ]] && continue
      ufw route allow in on "${WG_IF}" to "${net}" >/dev/null 2>&1 || true
    done
    ufw reload >/dev/null
  fi
}

configure_firewalld() {
  if command -v firewall-cmd >/dev/null 2>&1 && firewall-cmd --state >/dev/null 2>&1; then
    firewall-cmd --permanent --zone="${FIREWALLD_ZONE}" --add-port=${WG_PORT}/udp >/dev/null
    firewall-cmd --permanent --zone="${FIREWALLD_ZONE}" --add-masquerade >/dev/null
    firewall-cmd --permanent --direct --add-rule ipv4 filter FORWARD 0 \
      "iif ${WG_IF} oif ${WAN_IF} -s ${WG_NET} -j ACCEPT" >/dev/null 2>&1 || true
    firewall-cmd --permanent --direct --add-rule ipv4 filter FORWARD 0 \
      "oif ${WG_IF} -d ${WG_NET} -m conntrack --ctstate RELATED,ESTABLISHED -j ACCEPT" >/dev/null 2>&1 || true
    for net in "${EXTRA_NETS_ARRAY[@]}"; do
      [[ -z "${net}" ]] && continue
      firewall-cmd --permanent --direct --add-rule ipv4 filter FORWARD 0 \
        "iif ${WG_IF} -d ${net} -j ACCEPT" >/dev/null 2>&1 || true
    done
    firewall-cmd --reload >/dev/null
  fi
}

ensure_service() {
  systemctl enable "wg-quick@${WG_IF}" >/dev/null
  systemctl restart "wg-quick@${WG_IF}"
}

show_status() {
  echo "WireGuard routing configured with ${WG_IF} (${WG_ADDR}) via ${WAN_IF}"
  wg show "${WG_IF}" || true
  ip route show table main | grep "${WG_NET}" || true
}

main() {
  require_root
  ensure_sysctl
  backend=$(detect_backend)
  case "${backend}" in
    iptables) apply_iptables ;;
    nftables) apply_nftables ;;
  esac
  configure_ufw
  configure_firewalld
  ensure_service
  show_status
}

main "$@"