#!/usr/bin/env bash
# Sets up a point-to-point WireGuard tunnel between the two Swarm hosts. 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 is joined over this tunnel's private IPs instead.
#
# Run on EACH box, in either order. Each run generates that box's own WireGuard
# keypair (if missing) and prints its public key — paste it into the OTHER
# box's invocation. Two runs total, one per box:
#
#   box A (prod):  bash setup-wireguard.sh 10.10.0.1 10.10.0.2 <PROD_PUBLIC_IP>  <BETA_PUBLIC_IP>  <BETA_WG_PUBKEY_or_empty>
#   box B (beta):  bash setup-wireguard.sh 10.10.0.2 10.10.0.1 <BETA_PUBLIC_IP>  <PROD_PUBLIC_IP>  <PROD_WG_PUBKEY>
#
# First run on either box: leave the last argument empty, note the printed
# pubkey, then run the second box with that value, then re-run the first box
# once more with the second box's now-known pubkey. (A local key never
# changes across re-runs — only the peer section updates.)
set -euo pipefail

MY_WG_IP="${1:?usage: $0 <my_wg_ip> <peer_wg_ip> <my_public_ip> <peer_public_ip> [peer_wg_pubkey]}"
PEER_WG_IP="${2:?}"
MY_PUBLIC_IP="${3:?}"
PEER_PUBLIC_IP="${4:?}"
PEER_PUBKEY="${5:-}"

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

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 box'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 box's WireGuard public key (give this to the OTHER box's run):"
echo "    $MY_PUBKEY"
echo

if [ -z "$PEER_PUBKEY" ]; then
  echo "No peer public key supplied yet — writing a config with no [Peer] section."
  echo "Run this same command again once you have it (from the other box's output above)."
  cat > "$WG_DIR/wg0.conf" <<EOF
[Interface]
Address = ${MY_WG_IP}/24
PrivateKey = ${MY_PRIVKEY}
ListenPort = ${WG_PORT}
EOF
else
  cat > "$WG_DIR/wg0.conf" <<EOF
[Interface]
Address = ${MY_WG_IP}/24
PrivateKey = ${MY_PRIVKEY}
ListenPort = ${WG_PORT}

[Peer]
PublicKey = ${PEER_PUBKEY}
Endpoint = ${PEER_PUBLIC_IP}:${WG_PORT}
AllowedIPs = ${PEER_WG_IP}/32
PersistentKeepalive = 25
EOF
fi
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)
# Re-apply if the config changed on an already-up interface.
wg syncconf wg0 <(wg-quick strip wg0) 2>/dev/null || true

echo "==> Firewall: only let the OTHER box's public IP reach the WireGuard port"
if command -v ufw >/dev/null 2>&1; then
  ufw allow from "$PEER_PUBLIC_IP" to any port "$WG_PORT" proto udp comment "wireguard peer"
fi

echo
echo "wg0 status:"
wg show wg0 || true
echo
if [ -n "$PEER_PUBKEY" ]; then
  echo "==> Verify from here once BOTH boxes have run with each other's pubkey:"
  echo "    ping -c2 ${PEER_WG_IP}"
fi