#!/usr/bin/env bash
# Sets up this node's side of a full-mesh WireGuard tunnel across all of
# prod, beta, and the desktop (three nodes, not two — see
# docs/runbooks/implementation-handoff.md Track B step 2). Docker Swarm's
# control plane (2377/tcp) is TLS-encrypted by default, but the overlay data
# plane (VXLAN, 4789/udp) is NOT — and it should never face the public
# internet. Swarm joins over these private WireGuard IPs instead.
#
# A full mesh means every node peers directly with every other node (not a
# hub-and-spoke through one box) — with three nodes that's 2 [Peer] blocks per
# node. Run this on EACH of the three nodes, in any order, using a small peer
# list file so the script isn't hardcoded to any particular node count.
#
# --- Peer list format --------------------------------------------------
# A text file, one line per OTHER node (not including the one you're running
# on), each line: <wg_ip> <public_ip> <pubkey-or-dash>
#   10.10.0.1 169.58.46.181  <prod's pubkey, or - if not yet known>
#   10.10.0.2 75.119.132.91  <beta's pubkey, or ->
#   10.10.0.3 <desktop's public/reachable IP or a DDNS name>  <desktop's pubkey, or ->
#
# First pass on any node: peer pubkeys you don't have yet are "-". Run this
# script on all three nodes once (prints each node's own pubkey), fill those
# into everyone's peer-list files, then re-run on all three once more to
# actually establish the tunnels. Re-running is safe (idempotent) at any point.
#
#   bash setup-wireguard.sh <my_wg_ip> <peers_file>
set -euo pipefail

MY_WG_IP="${1:?usage: $0 <my_wg_ip> <peers_file>}"
PEERS_FILE="${2:?}"
WG_PORT="${WG_PORT:-51820}"
WG_DIR=/etc/wireguard

if [ "$(id -u)" -ne 0 ]; then
  echo "Run as root (or via the agent user's sudo)." >&2
  exit 1
fi
if [ ! -f "$PEERS_FILE" ]; then
  echo "Peer list not found: $PEERS_FILE (see this script's header for the format)" >&2
  exit 1
fi

echo "==> Installing WireGuard if needed"
command -v wg >/dev/null 2>&1 || { apt-get update -y -q && apt-get install -y -q wireguard; }

install -d -m 700 "$WG_DIR"
if [ ! -f "$WG_DIR/privatekey" ]; then
  echo "==> Generating this node's WireGuard keypair"
  umask 077
  wg genkey | tee "$WG_DIR/privatekey" | wg pubkey > "$WG_DIR/publickey"
fi
MY_PRIVKEY="$(cat "$WG_DIR/privatekey")"
MY_PUBKEY="$(cat "$WG_DIR/publickey")"

echo
echo "==> This node's WireGuard public key (put this in the OTHER two nodes'"
echo "    peer-list files, replacing the '-' for this node's line):"
echo "    $MY_PUBKEY"
echo

missing=0
{
  echo "[Interface]"
  echo "Address = ${MY_WG_IP}/24"
  echo "PrivateKey = ${MY_PRIVKEY}"
  echo "ListenPort = ${WG_PORT}"
  echo

  while read -r peer_wg_ip peer_public_ip peer_pubkey; do
    [ -z "${peer_wg_ip:-}" ] && continue
    case "$peer_wg_ip" in \#*) continue ;; esac
    if [ "$peer_pubkey" = "-" ] || [ -z "$peer_pubkey" ]; then
      echo "# SKIPPED (no pubkey yet) — peer at ${peer_public_ip}" >&2
      missing=$((missing + 1))
      continue
    fi
    echo "[Peer]"
    echo "PublicKey = ${peer_pubkey}"
    echo "Endpoint = ${peer_public_ip}:${WG_PORT}"
    echo "AllowedIPs = ${peer_wg_ip}/32"
    echo "PersistentKeepalive = 25"
    echo
  done < "$PEERS_FILE"
} > "$WG_DIR/wg0.conf"
chmod 600 "$WG_DIR/wg0.conf"

echo "==> Bringing up wg0"
systemctl enable --now wg-quick@wg0 2>/dev/null || (wg-quick down wg0 2>/dev/null; wg-quick up wg0)
wg syncconf wg0 <(wg-quick strip wg0) 2>/dev/null || true

echo "==> Firewall: only let listed peers' public IPs reach the WireGuard port"
if command -v ufw >/dev/null 2>&1; then
  while read -r _wg_ip peer_public_ip _pubkey; do
    [ -z "${peer_public_ip:-}" ] && continue
    case "$peer_public_ip" in \#*) continue ;; esac
    ufw allow from "$peer_public_ip" to any port "$WG_PORT" proto udp comment "wireguard peer"
  done < "$PEERS_FILE"
fi

echo
echo "wg0 status:"
wg show wg0 || true
echo
if [ "$missing" -gt 0 ]; then
  echo "NOTE: $missing peer(s) skipped (no pubkey yet in $PEERS_FILE). Fill them in"
  echo "and re-run on this node once all three nodes have printed their pubkey."
else
  echo "==> All peers configured. Verify once every node has run this:"
  while read -r peer_wg_ip _peer_public_ip peer_pubkey; do
    [ -z "${peer_wg_ip:-}" ] && continue
    case "$peer_wg_ip" in \#*) continue ;; esac
    [ "$peer_pubkey" = "-" ] && continue
    echo "    ping -c2 ${peer_wg_ip}"
  done < "$PEERS_FILE"
fi