#!/usr/bin/env bash
# Configure ufw firewall for this server stack:
#   - SSH (22)       : VPN clients only (10.13.13.0/24)
#   - SMTP (25)      : public (external mail servers must reach this)
#   - HTTP/HTTPS     : public (Traefik)
#   - WireGuard UDP  : public
#   - Mail client ports (465, 587, 143, 993, 4190): public by default
#     Set MAIL_VPN_ONLY=true to restrict these to VPN clients as well.
#
# Run this AFTER the WireGuard container is up and you have verified you can
# connect to the server via VPN (10.13.13.x). Once applied, SSH on port 22 is
# only reachable from within the WireGuard subnet.
#
# Usage:
#   sudo bash scripts/secure-ssh-vpn.sh [--dry-run] [--mail-vpn-only]
#
# To undo SSH restriction:
#   sudo ufw delete deny 22
#   sudo ufw delete allow from 10.13.13.0/24 to any port 22

set -euo pipefail

WG_SUBNET="10.13.13.0/24"
SSH_PORT="22"
WG_UDP_PORT="51820"
DRY_RUN=false
MAIL_VPN_ONLY=false

for arg in "$@"; do
  [[ "$arg" == "--dry-run" ]]      && DRY_RUN=true
  [[ "$arg" == "--mail-vpn-only" ]] && MAIL_VPN_ONLY=true
done

run() {
  if $DRY_RUN; then
    echo "[dry-run] $*"
  else
    "$@"
  fi
}

if [[ $EUID -ne 0 ]]; then
  echo "Error: run this script as root (sudo)." >&2
  exit 1
fi

echo "=== WireGuard SSH firewall setup ==="
echo "VPN subnet      : $WG_SUBNET"
echo "SSH port        : $SSH_PORT (VPN only)"
echo "WG UDP port     : $WG_UDP_PORT/udp"
echo "Mail VPN only   : $MAIL_VPN_ONLY"
$DRY_RUN && echo "(dry-run mode — no changes will be made)"
echo

# Ensure ufw is installed
if ! command -v ufw &>/dev/null; then
  echo "Installing ufw..."
  run apt-get install -y ufw
fi

# ── Web & VPN ──────────────────────────────────────────────────────────────
run ufw allow 80/tcp             comment 'HTTP (Traefik)'
run ufw allow 443/tcp            comment 'HTTPS (Traefik)'
run ufw allow 2424/tcp           comment 'GitLab SSH (Traefik)'
run ufw allow "$WG_UDP_PORT/udp" comment 'WireGuard VPN'

# ── SSH — VPN clients only ────────────────────────────────────────────────
# Deny rule is added first (gets higher rule number), then the VPN allow is
# inserted at position 1 so it always takes priority over the deny.
run ufw deny "$SSH_PORT/tcp"     comment 'Block public SSH'
run ufw insert 1 allow from "$WG_SUBNET" to any port "$SSH_PORT" comment 'SSH via WireGuard VPN'

# ── Mail (Mailcow) ────────────────────────────────────────────────────────
# Port 25 must always be public so other mail servers can deliver to you.
run ufw allow 25/tcp             comment 'SMTP (public — required for mail delivery)'

if $MAIL_VPN_ONLY; then
  # Restrict mail client access to VPN subnet only
  for port in 465 587 143 993 4190; do
    run ufw deny "$port/tcp" comment "Block public mail client port $port"
    run ufw insert 1 allow from "$WG_SUBNET" to any port "$port" comment "Mail client port $port via VPN"
  done
  echo "Mail client ports (465, 587, 143, 993, 4190) restricted to VPN."
else
  run ufw allow 465/tcp          comment 'SMTPS (Mailcow)'
  run ufw allow 587/tcp          comment 'SMTP Submission (Mailcow)'
  run ufw allow 143/tcp          comment 'IMAP (Mailcow)'
  run ufw allow 993/tcp          comment 'IMAPS (Mailcow)'
  run ufw allow 4190/tcp         comment 'Sieve (Mailcow)'
fi

# ── Enable / reload ufw ───────────────────────────────────────────────────
if ! ufw status | grep -q "Status: active"; then
  echo
  echo "WARNING: enabling ufw. Make sure you are connected via WireGuard"
  echo "         or you will lose your current SSH session."
  echo "Press Ctrl-C within 10 seconds to abort..."
  sleep 10
  run ufw --force enable
else
  run ufw reload
fi

echo
echo "Done. Current ufw rules:"
ufw status numbered
echo
echo "To verify WireGuard is running:"
echo "  docker compose -f /var/deploy/core/docker-compose.yml ps wireguard"
echo "  docker exec wireguard wg show"
echo
echo "Client configs are in: ./core/volumes/wireguard/config/peer_<name>/"