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

script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
# shellcheck source=scripts/cluster/common.sh
source "$script_dir/common.sh"

load_cluster_env
hosts=("$CLUSTER_NINGBO_HOST" "$CLUSTER_HONGKONG_HOST" "$CLUSTER_LOS_ANGELES_HOST")

for host in "${hosts[@]}"; do
  echo "[firewall] hardening K3s ports on $host"
  cluster_ssh "$host" 'bash -s' <<'REMOTE'
set -euo pipefail
[[ -d /sys/class/net/wg0 ]] || { echo 'wg0 is unavailable' >&2; exit 1; }
public_interface=$(ip -o route get 1.1.1.1 | awk '{ for (i = 1; i <= NF; i++) if ($i == "dev") { print $(i + 1); exit } }')
[[ $public_interface =~ ^[A-Za-z0-9._-]+$ && $public_interface != wg0 ]] || {
  echo 'cannot determine public interface' >&2
  exit 1
}
export DEBIAN_FRONTEND=noninteractive
if ! command -v netfilter-persistent >/dev/null 2>&1; then
  apt-get update -qq
  apt-get install -y -qq iptables-persistent
fi
ensure_rule() {
  if ! iptables -C INPUT "$@" 2>/dev/null; then
    iptables -I INPUT 1 "$@"
  fi
}
ensure_rule -i wg0 -s 10.77.0.0/24 -p tcp -m multiport \
  --dports 2379,2380,6443,10250 -m comment --comment easyai-k3s-private -j ACCEPT
ensure_rule -i wg0 -s 10.77.0.0/24 -p udp --dport 8472 \
  -m comment --comment easyai-flannel-private -j ACCEPT
ensure_rule -i wg0 -s 10.77.0.0/24 -p tcp --dport 30000:32767 \
  -m comment --comment easyai-nodeport-private -j ACCEPT
ensure_rule -i "$public_interface" -p udp --dport 51820 \
  -m comment --comment easyai-wireguard-public -j ACCEPT
ensure_rule -i "$public_interface" -p tcp -m multiport \
  --dports 2379,2380,6443,10250 -m comment --comment easyai-k3s-public-drop -j DROP
ensure_rule -i "$public_interface" -p udp --dport 8472 \
  -m comment --comment easyai-flannel-public-drop -j DROP
ensure_rule -i "$public_interface" -p tcp --dport 30000:32767 \
  -m comment --comment easyai-nodeport-public-drop -j DROP
netfilter-persistent save >/dev/null
REMOTE
done

echo 'node_network_hardening=PASS nodes=3'