#!/usr/bin/env bash
set -euo pipefail
ADDED_V4=0
ADDED_V6=0

trap 'echo "[gwiaprotection] ERROR at ${BASH_SOURCE[0]}:${LINENO}: ${BASH_COMMAND}" >&2' ERR

# gwiaprotection.sh
# Lock down a specific GWIA IP:PORT(S) so only MX hosts (plus internal nets/extra IPs) can connect.
# Usage:
#   /etc/scripts/gwiaprotection.sh --ip 192.168.10.250 --config
#   /etc/scripts/gwiaprotection.sh --ip 192.168.10.250 --status
#   /etc/scripts/gwiaprotection.sh --ip 192.168.10.250 --clear
#
# Loads config from /etc/scripts/gwiaprotection.<A.B.C.D>.conf (derived from --ip)
# Config keys:
#   GWIA_PORTS="2525 25"
#   DOMAINS=("eacbelcher.net")
#   INTERNAL_NETS="192.168.10.0/24"
#   EXTRA_IPS_V4=""
#   EXTRA_IPS_V6=""

die(){ echo "Error: $*" >&2; exit 1; }
log(){ logger -t "gwiaprotection" -- "$*"; echo "[gwiaprotection] $*"; }
need_root(){ [[ $EUID -eq 0 ]] || die "Run as root."; }

fw(){ firewall-cmd "$@"; }
fw_perm(){ firewall-cmd --permanent "$@"; }

# ---------------- args ----------------
ACTION=""; DST_IP=""
while [[ $# -gt 0 ]]; do
  case "$1" in
    --ip) shift; DST_IP="${1:-}";;
    --config|--status|--clear) ACTION="$1";;
    *) die "Unknown arg: $1";;
  esac
  shift || true
done

[[ -n "${DST_IP}" ]] || die "--ip <A.B.C.D> is required"
[[ "${DST_IP}" =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]] || die "Invalid IPv4: ${DST_IP}"
[[ -n "${ACTION}" ]] || die "Choose one: --config | --status | --clear"

DEBUG="${DEBUG:-0}"
if [[ "$DEBUG" == "1" ]]; then set -x; fi


CONF="/etc/scripts/gwiaprotection.${DST_IP}.conf"
[[ -f "${CONF}" ]] || die "Config not found: ${CONF}"
# shellcheck disable=SC1090
source "${CONF}"

# Normalize DOMAINS: allow either Bash array or space-separated string
if [[ -z "${DOMAINS[*]:-}" && -n "${DOMAINS:-}" ]]; then
  # DOMAINS provided as string; convert to array
  read -r -a DOMAINS <<<"$DOMAINS"
fi

# Fail if still empty
if [[ -z "${DOMAINS[*]:-}" ]]; then
  die "DOMAINS is empty in ${CONF} (use DOMAINS=(\"example.com\") or DOMAINS=\"example.com\")"
fi

# Also sanity check GWIA_PORTS
if [[ -z "${GWIA_PORTS:-}" ]]; then
  die "GWIA_PORTS is empty in ${CONF}"
fi

if [[ "$DEBUG" == "1" ]]; then
  echo "[DEBUG] Using CONF=${CONF}"
  echo "[DEBUG] DST_IP=${DST_IP}"
  echo "[DEBUG] GWIA_PORTS=${GWIA_PORTS}"
  echo -n "[DEBUG] DOMAINS=("; printf "%s " "${DOMAINS[@]}"; echo ")"
  echo "[DEBUG] INTERNAL_NETS=${INTERNAL_NETS}"
fi

[[ -n "${GWIA_PORTS:-}" ]] || die "GWIA_PORTS not set in ${CONF}"
[[ -n "${DOMAINS[*]:-}" ]] || die "DOMAINS[] not set in ${CONF}"
#INTERNAL_NETS="${INTERNAL_NETS:-}"
INTERNAL_NETS=""
EXTRA_IPS_V4="${EXTRA_IPS_V4:-}"
EXTRA_IPS_V6="${EXTRA_IPS_V6:-}"

SAFE_IP="${DST_IP//./_}"
# Short names to satisfy ipset 31-char limit (e.g. gwp_192_168_10_250_v4)
IPSET_V4="gwp_${SAFE_IP}_v4"
IPSET_V6="gwp_${SAFE_IP}_v6"

ensure_firewalld(){ systemctl is-active --quiet firewalld || systemctl start firewalld; }
zone(){ fw --get-default-zone; }

purge_legacy_test_rules(){
  local z; z="$(zone)"
  for p in ${GWIA_PORTS:-25}; do
    fw_perm --zone="$z" \
      --remove-rich-rule="rule priority=\"-200\" family=\"ipv4\" destination address=\"${DST_IP}\" source ipset=\"gwiatest_v4\" port port=\"${p}\" protocol=\"tcp\" accept" 2>/dev/null || true
  done
}

ensure_ipsets(){
  # Just ensure permanent ipsets exist, then reload once. No deletes here.
  if ! fw_perm --get-ipsets | grep -qx "${IPSET_V4}"; then
    fw_perm --new-ipset="${IPSET_V4}" --type=hash:ip --option=family=inet
  fi
  # Create v6 set only if we’ll ever need it later; still OK to create now:
  if ! fw_perm --get-ipsets | grep -qx "${IPSET_V6}"; then
    fw_perm --new-ipset="${IPSET_V6}" --type=hash:ip --option=family=inet6
  fi
  fw --reload
}

resolve_mx_ips(){
  local domain="$1"
  local mx h
  mx=$(dig +short MX "$domain" 2>/dev/null | awk '{print $2}' | sed 's/\.$//') || true
  [[ -z "$mx" ]] && return 0
  for h in $mx; do
    dig +short A "$h" 2>/dev/null || true
    dig +short AAAA "$h" 2>/dev/null || true
  done
}

populate_ipsets(){
  ADDED_V4=0; ADDED_V6=0

  # Start from clean slate (both runtime and permanent)
  if fw --ipset="${IPSET_V4}" --get-entries >/dev/null 2>&1 || true; then
    while read -r e; do [[ -n "$e" ]] && fw --ipset="${IPSET_V4}" --remove-entry="$e"; done < <(fw --ipset="${IPSET_V4}" --get-entries || true)
  fi
  if fw_perm --ipset="${IPSET_V4}" --get-entries >/dev/null 2>&1 || true; then
    while read -r e; do [[ -n "$e" ]] && fw_perm --ipset="${IPSET_V4}" --remove-entry="$e"; done < <(fw_perm --ipset="${IPSET_V4}" --get-entries || true)
  fi
  if fw --ipset="${IPSET_V6}" --get-entries >/dev/null 2>&1 || true; then
    while read -r e; do [[ -n "$e" ]] && fw --ipset="${IPSET_V6}" --remove-entry="$e"; done < <(fw --ipset="${IPSET_V6}" --get-entries || true)
  fi
  if fw_perm --ipset="${IPSET_V6}" --get-entries >/dev/null 2>&1 || true; then
    while read -r e; do [[ -n "$e" ]] && fw_perm --ipset="${IPSET_V6}" --remove-entry="$e"; done < <(fw_perm --ipset="${IPSET_V6}" --get-entries || true)
  fi

  # Populate from MX
  if ! command -v dig >/dev/null 2>&1; then
    echo "[gwiaprotection] WARNING: 'dig' not found; skipping MX population" >&2
  fi

  local d ip
  for d in "${DOMAINS[@]}"; do
    while read -r ip; do
      [[ -z "$ip" ]] && continue
      if [[ "$ip" == *:* ]]; then
        fw --ipset="${IPSET_V6}" --add-entry="$ip" 2>/dev/null || true
        fw_perm --ipset="${IPSET_V6}" --add-entry="$ip" 2>/dev/null || true
        ((ADDED_V6++))
      else
        fw --ipset="${IPSET_V4}" --add-entry="$ip" 2>/dev/null || true
        fw_perm --ipset="${IPSET_V4}" --add-entry="$ip" 2>/dev/null || true
        ((ADDED_V4++))
      fi
    done < <(resolve_mx_ips "$d")
  done

  # Extras
  for ip in $EXTRA_IPS_V4; do
    fw --ipset="${IPSET_V4}" --add-entry="$ip" 2>/dev/null || true
    fw_perm --ipset="${IPSET_V4}" --add-entry="$ip" 2>/dev/null || true
    ((ADDED_V4++))
  done
  for ip in $EXTRA_IPS_V6; do
    fw --ipset="${IPSET_V6}" --add-entry="$ip" 2>/dev/null || true
    fw_perm --ipset="${IPSET_V6}" --add-entry="$ip" 2>/dev/null || true
    ((ADDED_V6++))
  done

  [[ "$DEBUG" == "1" ]] && echo "[DEBUG] populate_ipsets: added_v4=${ADDED_V4} added_v6=${ADDED_V6}"
}


remove_rules(){
  local z p net; z="$(zone)"
  purge_legacy_test_rules
  for p in ${GWIA_PORTS:-25}; do
    for net in ${INTERNAL_NETS:-}; do
      fw_perm --zone="$z" --remove-rich-rule="rule priority=\"-200\" family=\"ipv4\" destination address=\"${DST_IP}\" source address=\"${net}\" port port=\"${p}\" protocol=\"tcp\" accept" 2>/dev/null || true
    done
    fw_perm --zone="$z" --remove-rich-rule="rule priority=\"-200\" family=\"ipv4\" destination address=\"${DST_IP}\" source ipset=\"${IPSET_V4}\" port port=\"${p}\" protocol=\"tcp\" accept" 2>/dev/null || true
    fw_perm --zone="$z" --remove-rich-rule="rule priority=\"-200\" family=\"ipv6\" destination address=\"${DST_IP}\" source ipset=\"${IPSET_V6}\" port port=\"${p}\" protocol=\"tcp\" accept" 2>/dev/null || true
    fw_perm --zone="$z" --remove-rich-rule="rule priority=\"100\" family=\"ipv4\" destination address=\"${DST_IP}\" port port=\"${p}\" protocol=\"tcp\" drop" 2>/dev/null || true
  done
  fw --reload || true
}

add_rules(){
echo "[DEBUG] add_rules(): entering"
  local z p net; z="$(zone)"
  [[ -z "$z" ]] && z="public"
  [[ "$DEBUG" == "1" ]] && echo "[DEBUG] add_rules: zone=$z dst=$DST_IP ports=$GWIA_PORTS"

  for p in $GWIA_PORTS; do
    for net in $INTERNAL_NETS; do
echo "[DEBUG] add_rules(): adding LAN-accept for $DST_IP:$p from $net"
echo "[DEBUG] add_rules(): adding IPSET-accept for $DST_IP:$p (v4 set ${IPSET_V4})"
echo "[DEBUG] add_rules(): adding DROP for $DST_IP:$p"
      fw_perm --zone="$z" --add-rich-rule="rule priority=\"-200\" family=\"ipv4\" destination address=\"${DST_IP}\" source address=\"${net}\" port port=\"${p}\" protocol=\"tcp\" accept"
    done
#echo "[DEBUG] add_rules(): adding LAN-accept for $DST_IP:$p from $net"
echo "[DEBUG] add_rules(): adding IPSET-accept for $DST_IP:$p (v4 set ${IPSET_V4})"
echo "[DEBUG] add_rules(): adding DROP for $DST_IP:$p"
    fw_perm --zone="$z" --add-rich-rule="rule priority=\"-200\" family=\"ipv4\" destination address=\"${DST_IP}\" source ipset=\"${IPSET_V4}\" port port=\"${p}\" protocol=\"tcp\" accept"
    if [[ $ADDED_V6 -gt 0 ]]; then
echo "[DEBUG] add_rules(): adding LAN-accept for $DST_IP:$p from $net"
echo "[DEBUG] add_rules(): adding IPSET-accept for $DST_IP:$p (v4 set ${IPSET_V4})"
echo "[DEBUG] add_rules(): adding DROP for $DST_IP:$p"
      fw_perm --zone="$z" --add-rich-rule="rule priority=\"-200\" family=\"ipv6\" destination address=\"${DST_IP}\" source ipset=\"${IPSET_V6}\" port port=\"${p}\" protocol=\"tcp\" accept"
    fi
#echo "[DEBUG] add_rules(): adding LAN-accept for $DST_IP:$p from $net"
echo "[DEBUG] add_rules(): adding IPSET-accept for $DST_IP:$p (v4 set ${IPSET_V4})"
echo "[DEBUG] add_rules(): adding DROP for $DST_IP:$p"
    fw_perm --zone="$z" --add-rich-rule="rule priority=\"100\" family=\"ipv4\" destination address=\"${DST_IP}\" port port=\"${p}\" protocol=\"tcp\" drop"
  done
  fw --reload
}

status(){
  local z; z="$(zone)"
  echo "=== Zone: ${z} | DST_IP: ${DST_IP} ==="
  echo "Rich rules for ${DST_IP}:"
  fw --zone="${z}" --list-rich-rules | grep "destination address=\"${DST_IP}\"" || echo "(none)"
  echo
  echo "IPv4 ipset (${IPSET_V4}) RUNTIME:";   fw --ipset="${IPSET_V4}" --get-entries 2>/dev/null || echo "(none)"
  echo "IPv4 ipset (${IPSET_V4}) PERMANENT:"; fw_perm --ipset="${IPSET_V4}" --get-entries 2>/dev/null || echo "(none)"
  echo "IPv6 ipset (${IPSET_V6}) RUNTIME:";   fw --ipset="${IPSET_V6}" --get-entries 2>/dev/null || echo "(none)"
  echo "IPv6 ipset (${IPSET_V6}) PERMANENT:"; fw_perm --ipset="${IPSET_V6}" --get-entries 2>/dev/null || echo "(none)"
}

clear_all(){
  ensure_firewalld
  purge_legacy_test_rules
  remove_rules
  fw_perm --delete-ipset="${IPSET_V4}" 2>/dev/null || true
  fw_perm --delete-ipset="${IPSET_V6}" 2>/dev/null || true
  fw --reload
  log "Cleared rules + ipsets for ${DST_IP}"
}




do_config(){
  need_root; ensure_firewalld
  purge_legacy_test_rules

  echo "[DEBUG] do_config: remove_rules"
  remove_rules || true

  echo "[DEBUG] do_config: ensure_ipsets"
  ensure_ipsets || true

  echo "[DEBUG] do_config: populate_ipsets"
  populate_ipsets || true

  echo "[DEBUG] do_config: add_rules"
  add_rules || true

  log "Configured ${DST_IP} for domains: ${DOMAINS[*]} on ports: ${GWIA_PORTS}"
}
case "${ACTION}" in
  --config) do_config ;;
  --status) status ;;
  --clear)  clear_all ;;
esac