#!/usr/bin/env bash
# =============================================================================
#  CAST Imaging port firewall for Linux (Docker / Podman)
#
#  Port table (see "Network Rules - Restricting CAST Imaging Ports to the Gateway"):
#    PUBLIC (any machine)             : Gateway 8090 (PORT_GATEWAY)
#    RESTRICTED (this server + other  : Control panel 8098, Console 8091,
#      Imaging servers only)            Auth service 8092, Keycloak/SSO 8096,
#                                       Message broker 2381, Analysis node 8089,
#                                       Dashboards 8097, Viewer 9000/9001/8083/8082/8070,
#                                       Neo4j Bolt 7687, + every other PORT key in the config,
#                                       + embedded PostgreSQL DB_PORT (2285) when DB_MODE=embedded
#    OUT OF SCOPE (not touched)       : external PostgreSQL (DB_MODE=external) - its port is on
#                                       another machine
#
#  Other Imaging servers are given as IPs, CIDRs or host names, e.g.
#    192.168.1.2, 198.168.2.5, abc-server1, abc-server2.corp.cast.com
#  either in IMAGING_SERVERS below, with -n, and/or automatically from the
#  *_HOSTNAME values of the config. Host names are resolved when the rules are
#  applied - re-run "apply" if a server's IP changes.
#
#  How: a dedicated nftables table hooked in PREROUTING *before* Docker/Podman
#  NAT, so published container ports cannot bypass it (they bypass ufw/firewalld).
#  Local access (loopback) and container-to-container traffic are never filtered.
#
#  Usage:
#    sudo ./cast-imaging-firewall.sh dryrun [options]   show rules, change nothing
#    sudo ./cast-imaging-firewall.sh apply  [options]   apply + persist at boot
#    sudo ./cast-imaging-firewall.sh remove             revert completely
#    sudo ./cast-imaging-firewall.sh status             show active rules + drop counter
#
#  Options:
#    -c FILE    configuration.conf or compose .env. Default: ./configuration.conf,
#               then configuration.conf next to this script.
#    -n LIST    other Imaging servers (IP, IP/CIDR or host name), comma separated,
#               e.g. -n "192.168.1.2,198.168.2.5,abc-server1,abc-server2.corp.cast.com"
#    -a LIST    extra config KEYS to open to any machine, comma separated
#               (e.g. -a PORT_CONTROL_PANEL)
#    -p LIST    extra ports to restrict (not in the config), comma separated
#    -H         do not derive Imaging servers from *_HOSTNAME values of the config
#    -D         also restrict DB_PORT even when DB_MODE=external
#    -N         do not persist (no systemd unit; rules lost at reboot)
# =============================================================================
set -euo pipefail

# -----------------------------------------------------------------------------
# Other CAST Imaging servers allowed to reach the restricted ports.
# Edit this list (IPs, IP/CIDR or host names), or use -n on the command line.
# -----------------------------------------------------------------------------
IMAGING_SERVERS=""
# IMAGING_SERVERS="192.168.1.2, 198.168.2.5, abc-server1, abc-server2.corp.cast.com"

TABLE="cast_imaging"
RULES_DIR="/etc/cast-imaging"
RULES_FILE="$RULES_DIR/firewall.nft"
UNIT_FILE="/etc/systemd/system/cast-imaging-firewall.service"

# PUBLIC - reachable from any machine. Role -> config keys / default port if no key is found.
declare -A PUBLIC_ROLE_KEYS=(    [gateway]="IMAGING_SERVICES_PORT_GATEWAY PORT_GATEWAY" )
declare -A PUBLIC_ROLE_DEFAULT=( [gateway]=8090 )

# RESTRICTED - this server + other Imaging servers only. Always restricted, even when
# missing from the config (defaults of the port table). Role -> config keys / default port.
declare -A RESTRICTED_ROLE_KEYS=(
  [control-panel]="IMAGING_SERVICES_PORT_CONTROL_PANEL PORT_CONTROL_PANEL"
  [console]="IMAGING_SERVICES_PORT_CONSOLE PORT_CONSOLE"
  [auth-service]="IMAGING_SERVICES_PORT_AUTH_SERVICE PORT_AUTH_SERVICE"
  [keycloak-sso]="IMAGING_SERVICES_PORT_SSO PORT_SSO PORT_KEYCLOAK"
  [message-broker]="BROKER_MQ_PORT"
  [analysis-node]="IMAGING_SERVICES_PORT_ANALYSIS_NODE PORT_ANALYSIS_NODE"
  [dashboards]="IMAGING_SERVICES_PORT_DASHBOARDS PORT_DASHBOARDS"
  [viewer-server]="IMAGING_VIEWER_SERVER_PORT"
  [viewer-etl]="IMAGING_VIEWER_ETL_PORT"
  [viewer-frontend]="IMAGING_VIEWER_FRONTEND_PORT"
  [viewer-ai]="IMAGING_VIEWER_AI_PORT"
  [viewer-api]="IMAGING_VIEWER_API_PORT"
  [neo4j-bolt]="NEO4J_BOLT_PORT"
)
declare -A RESTRICTED_ROLE_DEFAULT=(
  [control-panel]=8098 [console]=8091 [auth-service]=8092 [keycloak-sso]=8096
  [message-broker]=2381 [analysis-node]=8089 [dashboards]=8097
  [viewer-server]=9000 [viewer-etl]=9001 [viewer-frontend]=8083 [viewer-ai]=8082 [viewer-api]=8070
  [neo4j-bolt]=7687
)

# PostgreSQL: DB_MODE=embedded (installer default) -> DB_PORT restricted like the other ports;
#             DB_MODE=external                    -> DB_PORT out of scope (not touched)
DB_KEYS=( DB_PORT )
DB_DEFAULT_PORT=2285
FORCE_DB=0
EXCLUDED_KEYS=()
# Container bridge / veth interfaces: container-to-container traffic is never filtered
CONTAINER_IFACES=( "docker0" "podman*" "br-*" "cni*" "veth*" )

CONF_FILE=""
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
EXTRA_NODES=""; EXTRA_PUBLIC_KEYS=""; EXTRA_PORTS=""
NODES_FROM_HOSTNAMES=1; PERSIST=1

c_cyan=$'\e[36m'; c_yel=$'\e[33m'; c_grn=$'\e[32m'; c_red=$'\e[31m'; c_dim=$'\e[2m'; c_off=$'\e[0m'
info() { echo "${c_cyan}$*${c_off}"; }
warn() { echo "${c_yel}WARNING: $*${c_off}" >&2; }
die()  { echo "${c_red}ERROR: $*${c_off}" >&2; exit 1; }
usage() { awk 'NR>2 && /^# =+$/ {exit} NR>2' "$0" | sed 's/^# \{0,1\}//'; exit "${1:-0}"; }

ACTION="${1:-}"; [[ -z "$ACTION" || "$ACTION" == -h || "$ACTION" == --help ]] && usage 0
shift
while getopts ":c:n:a:p:HDN" opt; do
  case $opt in
    c) CONF_FILE="$OPTARG" ;;
    n) EXTRA_NODES="$OPTARG" ;;
    a) EXTRA_PUBLIC_KEYS="$OPTARG" ;;
    p) EXTRA_PORTS="$OPTARG" ;;
    H) NODES_FROM_HOSTNAMES=0 ;;
    D) FORCE_DB=1 ;;
    N) PERSIST=0 ;;
    *) usage 1 ;;
  esac
done

[[ $EUID -eq 0 ]] || die "run as root (sudo)."
NFT="$(command -v nft || true)"
[[ -n "$NFT" ]] || die "nft (nftables) not found - install it: dnf install nftables / apt install nftables"

# -----------------------------------------------------------------------------
# remove / status
# -----------------------------------------------------------------------------
if [[ "$ACTION" == "remove" ]]; then
  if [[ -f "$UNIT_FILE" ]]; then
    systemctl disable --now cast-imaging-firewall.service 2>/dev/null || true
    rm -f "$UNIT_FILE"; systemctl daemon-reload
    echo "Removed systemd unit."
  fi
  if "$NFT" list table inet "$TABLE" >/dev/null 2>&1; then
    "$NFT" delete table inet "$TABLE"; echo "Removed nftables table inet $TABLE."
  else
    echo "nftables table inet $TABLE not present."
  fi
  rm -f "$RULES_FILE"
  echo "${c_grn}Done - all CAST Imaging firewall rules removed.${c_off}"
  exit 0
fi

if [[ "$ACTION" == "status" ]]; then
  if "$NFT" list table inet "$TABLE" 2>/dev/null; then
    state="$(systemctl is-enabled cast-imaging-firewall.service 2>/dev/null || true)"
    echo; echo "Persistence (systemd): ${state:-not installed}"
  else
    echo "Not applied (no table inet $TABLE)."
  fi
  exit 0
fi

[[ "$ACTION" == "dryrun" || "$ACTION" == "apply" ]] || usage 1

# -----------------------------------------------------------------------------
# 1. Read configuration.conf
# -----------------------------------------------------------------------------
if [[ -z "$CONF_FILE" ]]; then
  for cand in "./configuration.conf" "$SCRIPT_DIR/configuration.conf"; do
    if [[ -f "$cand" ]]; then CONF_FILE="$cand"; break; fi
  done
  [[ -n "$CONF_FILE" ]] || die "configuration.conf not found in $(pwd) or $SCRIPT_DIR - pass it with -c <file>"
  echo "${c_dim}Using default config: $(readlink -f "$CONF_FILE")${c_off}"
fi
[[ -f "$CONF_FILE" ]] || die "config file not found: $CONF_FILE"

declare -A CONF; KEYS=()
while IFS= read -r line || [[ -n "$line" ]]; do
  line="${line%$'\r'}"                                    # CRLF files
  [[ "$line" =~ ^[[:space:]]*(#|$) ]] && continue
  if [[ "$line" =~ ^[[:space:]]*(export[[:space:]]+)?([A-Za-z0-9_]+)[[:space:]]*=(.*)$ ]]; then
    k="${BASH_REMATCH[2]}"; v="${BASH_REMATCH[3]}"
    v="${v#"${v%%[![:space:]]*}"}"; v="${v%"${v##*[![:space:]]}"}"   # trim
    v="${v%\"}"; v="${v#\"}"; v="${v%\'}"; v="${v#\'}"                # unquote
    [[ -v "CONF[$k]" ]] || KEYS+=("$k")
    CONF[$k]="$v"
  fi
done < "$CONF_FILE"

DB_MODE="$(tr '[:upper:]' '[:lower:]' <<< "${CONF[DB_MODE]:-embedded}")"
case "$DB_MODE" in
  embedded) ;;
  external) [[ $FORCE_DB -eq 1 ]] || EXCLUDED_KEYS=( "${DB_KEYS[@]}" ) ;;
  *) die "unknown DB_MODE='${CONF[DB_MODE]}' in $CONF_FILE (expected embedded or external)" ;;
esac
echo "${c_dim}Database mode: $DB_MODE$( [[ ${CONF[DB_MODE]+x} ]] || echo ' (DB_MODE not set - installer default)')${c_off}"

in_list() { local x="$1"; shift; local e; for e in "$@"; do [[ "$e" == "$x" ]] && return 0; done; return 1; }
split_list() { tr ',; \t' '\n\n\n\n' <<< "$1" | grep -v '^$' || true; }   # "a, b;c d" -> lines

PUBLIC_KEYS=()
for r in "${!PUBLIC_ROLE_KEYS[@]}"; do for k in ${PUBLIC_ROLE_KEYS[$r]}; do PUBLIC_KEYS+=("$k"); done; done
mapfile -t tmp < <(split_list "$EXTRA_PUBLIC_KEYS"); PUBLIC_KEYS+=("${tmp[@]+"${tmp[@]}"}")

PUBLIC_PORTS=(); RESTRICTED_PORTS=(); REPORT=()
for k in "${KEYS[@]+"${KEYS[@]}"}"; do
  v="${CONF[$k]}"
  [[ "$k" =~ (^|_)PORT(_|$) && "$v" =~ ^[0-9]+$ ]] || continue
  if in_list "$k" "${EXCLUDED_KEYS[@]+"${EXCLUDED_KEYS[@]}"}"; then REPORT+=("$k|$v|out of scope (external DB)"); continue; fi
  if in_list "$k" "${PUBLIC_KEYS[@]}"; then PUBLIC_PORTS+=("$v"); REPORT+=("$k|$v|PUBLIC")
  else RESTRICTED_PORTS+=("$v"); REPORT+=("$k|$v|RESTRICTED"); fi
done

role_has_key() {   # $1 = space-separated keys; true if one of them has a numeric value
  local k; for k in $1; do [[ "${CONF[$k]:-}" =~ ^[0-9]+$ ]] && return 0; done; return 1
}
for r in "${!PUBLIC_ROLE_KEYS[@]}"; do
  if ! role_has_key "${PUBLIC_ROLE_KEYS[$r]}"; then
    warn "$r port not set in config - using default ${PUBLIC_ROLE_DEFAULT[$r]}."
    PUBLIC_PORTS+=("${PUBLIC_ROLE_DEFAULT[$r]}"); REPORT+=("($r default)|${PUBLIC_ROLE_DEFAULT[$r]}|PUBLIC")
  fi
done
for r in "${!RESTRICTED_ROLE_KEYS[@]}"; do
  if ! role_has_key "${RESTRICTED_ROLE_KEYS[$r]}"; then
    RESTRICTED_PORTS+=("${RESTRICTED_ROLE_DEFAULT[$r]}"); REPORT+=("($r default)|${RESTRICTED_ROLE_DEFAULT[$r]}|RESTRICTED")
  fi
done

if [[ $DB_MODE == embedded ]] && ! role_has_key "${DB_KEYS[*]}"; then
  RESTRICTED_PORTS+=("$DB_DEFAULT_PORT"); REPORT+=("(embedded db default)|$DB_DEFAULT_PORT|RESTRICTED")
fi

mapfile -t tmp < <(split_list "$EXTRA_PORTS")
for p in "${tmp[@]+"${tmp[@]}"}"; do
  [[ "$p" =~ ^[0-9]+$ ]] || die "invalid port '$p' in -p"
  RESTRICTED_PORTS+=("$p"); REPORT+=("(-p extra)|$p|RESTRICTED")
done

uniq_sorted() { printf '%s\n' "$@" | grep -E '^[0-9]+$' | sort -n -u || true; }
mapfile -t PUBLIC_PORTS     < <(uniq_sorted "${PUBLIC_PORTS[@]+"${PUBLIC_PORTS[@]}"}")
mapfile -t RESTRICTED_PORTS < <(uniq_sorted "${RESTRICTED_PORTS[@]+"${RESTRICTED_PORTS[@]}"}")

for p in "${PUBLIC_PORTS[@]}" "${RESTRICTED_PORTS[@]}"; do
  (( p >= 1 && p <= 65535 )) || die "invalid port $p"
done

# A public port is never restricted
FINAL_RESTRICTED=()
for p in "${RESTRICTED_PORTS[@]}"; do
  if in_list "$p" "${PUBLIC_PORTS[@]}"; then warn "port $p is both public and restricted - kept public."; else FINAL_RESTRICTED+=("$p"); fi
done
[[ ${#FINAL_RESTRICTED[@]} -gt 0 ]] || die "no ports to restrict found in $CONF_FILE"

echo; info "Ports read from $CONF_FILE:"
printf '  %-40s %-6s %s\n' "KEY" "PORT" "ACCESS"
for r in "${REPORT[@]}"; do IFS='|' read -r k p a <<< "$r"; printf '  %-40s %-6s %s\n' "$k" "$p" "$a"; done

# -----------------------------------------------------------------------------
# 2. Other Imaging servers: IMAGING_SERVERS + -n + *_HOSTNAME values of the config
#    Accepts IPv4, IPv4/CIDR or host names (resolved to IPv4 now).
# -----------------------------------------------------------------------------
mapfile -t LOCAL_IPS < <( { hostname -I 2>/dev/null | tr ' ' '\n'; ip -o addr show 2>/dev/null | awk '{print $4}' | cut -d/ -f1; } | grep -v '^$' | sort -u )
LOCAL_NAMES=( localhost "$(hostname)" "$(hostname -f 2>/dev/null || hostname)" "$(hostname -s 2>/dev/null || hostname)" )

NODES=()
add_node() {   # $1 = entry, $2 = origin label, $3 = 1 if a failed resolution is fatal
  local n="$1" origin="$2" strict="$3" ip ips
  if [[ "$n" =~ ^[0-9]{1,3}(\.[0-9]{1,3}){3}(/[0-9]{1,2})?$ ]]; then
    NODES+=("$n"); echo "${c_dim}  server $n ($origin)${c_off}"; return
  fi
  [[ "$n" =~ ^[A-Za-z0-9]([A-Za-z0-9.-]*[A-Za-z0-9])?$ ]] || die "invalid Imaging server '$n' ($origin) - use IP, IP/CIDR or host name"
  mapfile -t ips < <(getent ahostsv4 "$n" 2>/dev/null | awk '{print $1}' | sort -u)
  if [[ ${#ips[@]} -eq 0 ]]; then
    if [[ "$strict" -eq 1 ]]; then die "cannot resolve Imaging server '$n' ($origin) - fix DNS/hosts or give its IP"; fi
    warn "cannot resolve $n ($origin) - add its IP with -n"; return
  fi
  for ip in "${ips[@]}"; do
    [[ "$ip" == 127.* ]] && continue
    in_list "$ip" "${LOCAL_IPS[@]+"${LOCAL_IPS[@]}"}" && continue
    NODES+=("$ip"); echo "${c_dim}  server $n -> $ip ($origin)${c_off}"
  done
}

echo; info "Other Imaging servers:"
mapfile -t tmp < <(split_list "$IMAGING_SERVERS")
for n in "${tmp[@]+"${tmp[@]}"}"; do add_node "$n" "IMAGING_SERVERS" 1; done
mapfile -t tmp < <(split_list "$EXTRA_NODES")
for n in "${tmp[@]+"${tmp[@]}"}"; do add_node "$n" "-n" 1; done
if [[ $NODES_FROM_HOSTNAMES -eq 1 ]]; then
  for k in "${KEYS[@]+"${KEYS[@]}"}"; do
    [[ "$k" =~ _HOSTNAME$ ]] || continue
    h="${CONF[$k]}"; [[ -z "$h" ]] && continue
    in_list "$h" "${LOCAL_NAMES[@]}" && continue
    add_node "$h" "$k" 0
  done
fi
mapfile -t NODES < <(printf '%s\n' "${NODES[@]+"${NODES[@]}"}" | grep -v '^$' | sort -u || true)
[[ ${#NODES[@]} -gt 0 ]] || echo "  none (standalone)"

echo
info "Summary:"
echo "  Public (any machine)          : TCP $(IFS=,; echo "${PUBLIC_PORTS[*]}")"
echo "  Restricted (local + servers)  : TCP $(IFS=,; echo "${FINAL_RESTRICTED[*]}")"
echo "  Allowed Imaging servers       : $( [[ ${#NODES[@]} -gt 0 ]] && (IFS=,; echo "${NODES[*]}") || echo 'none (standalone - restricted ports reachable only locally)')"

# -----------------------------------------------------------------------------
# 3. Published container ports not covered by the config (report only)
# -----------------------------------------------------------------------------
published_ports() {
  local eng; for eng in docker podman; do
    command -v "$eng" >/dev/null 2>&1 || continue
    "$eng" ps --format '{{.Ports}}' 2>/dev/null | tr ',' '\n' | sed 's/^ *//' |
      grep -- '->' | grep -v '^127\.0\.0\.1:' | grep -v '^\[::1\]:' |
      sed -E 's/->.*//; s/.*://' | grep -E '^[0-9]+(-[0-9]+)?$' || true
  done | sort -n -u
}
mapfile -t PUB < <(published_ports)
UNCOVERED=()
for p in "${PUB[@]+"${PUB[@]}"}"; do
  [[ "$p" =~ ^[0-9]+$ ]] || continue
  in_list "$p" "${PUBLIC_PORTS[@]}" "${FINAL_RESTRICTED[@]}" || UNCOVERED+=("$p")
done
if [[ ${#PUB[@]} -gt 0 ]]; then
  echo "  Published container ports     : $(IFS=,; echo "${PUB[*]}")"
  [[ ${#UNCOVERED[@]} -gt 0 ]] && warn "published on the network but not in the config (NOT restricted): $(IFS=,; echo "${UNCOVERED[*]}") - add with -p if internal."
else
  echo "${c_dim}  (no published ports found via docker/podman ps - rootless containers of other users are not listed)${c_off}"
fi

# -----------------------------------------------------------------------------
# 4. Build the nftables ruleset
# -----------------------------------------------------------------------------
join_by() { local IFS="$1"; shift; echo "$*"; }
NFT_RULES="$(cat <<NFTEOF
# CAST Imaging port firewall - generated by cast-imaging-firewall.sh on $(date '+%F %T')
# Source config: $(readlink -f "$CONF_FILE")
# Public (any machine): $(join_by ',' "${PUBLIC_PORTS[@]}")
table inet $TABLE
delete table inet $TABLE

table inet $TABLE {
  # Restricted ports: this server + other Imaging servers only
  set internal_ports {
    type inet_service
    elements = { $(join_by ', ' "${FINAL_RESTRICTED[@]}") }
  }

  # Other Imaging servers
  set nodes4 {
    type ipv4_addr
    flags interval
$( [[ ${#NODES[@]} -gt 0 ]] && echo "    elements = { $(join_by ', ' "${NODES[@]}") }" )
  }

  # priority -150: runs BEFORE Docker/Podman DNAT (dstnat = -100), sees original ports
  chain prerouting {
    type filter hook prerouting priority -150; policy accept;

    iifname "lo" accept
$(for i in "${CONTAINER_IFACES[@]}"; do echo "    iifname \"$i\" accept"; done)
    fib daddr type != local accept

    tcp dport @internal_ports ip saddr @nodes4 accept
    tcp dport @internal_ports limit rate 10/minute log prefix "cast-imaging-drop: "
    tcp dport @internal_ports counter drop
  }
}
NFTEOF
)"

echo; info "nftables ruleset:"
echo "$NFT_RULES" | sed 's/^/  /'

TMP_RULES="$(mktemp)"; trap 'rm -f "$TMP_RULES"' EXIT
echo "$NFT_RULES" > "$TMP_RULES"
"$NFT" -c -f "$TMP_RULES" || die "nftables syntax check failed"
echo "${c_grn}  Syntax check OK.${c_off}"

if command -v firewall-cmd >/dev/null 2>&1 && firewall-cmd --state >/dev/null 2>&1; then
  for p in "${PUBLIC_PORTS[@]}"; do
    firewall-cmd --query-port="$p/tcp" >/dev/null 2>&1 || \
      echo "${c_dim}  note: firewalld does not open $p/tcp - only matters for host-network containers (firewall-cmd --permanent --add-port=$p/tcp)${c_off}"
  done
fi

if [[ "$ACTION" == "dryrun" ]]; then
  echo; echo "${c_grn}DryRun - nothing changed.${c_off}"
  exit 0
fi

# -----------------------------------------------------------------------------
# 5. Apply + persist
# -----------------------------------------------------------------------------
mkdir -p "$RULES_DIR"
cp "$TMP_RULES" "$RULES_FILE"; chmod 0644 "$RULES_FILE"
"$NFT" -f "$RULES_FILE"
echo; echo "${c_grn}Applied: table inet $TABLE${c_off}"

if [[ $PERSIST -eq 1 ]]; then
  cat > "$UNIT_FILE" <<UNITEOF
[Unit]
Description=CAST Imaging port firewall
Wants=network-pre.target
After=network-pre.target nftables.service firewalld.service docker.service podman.service
Before=network-online.target

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=$NFT -f $RULES_FILE
ExecStop=-$NFT delete table inet $TABLE

[Install]
WantedBy=multi-user.target
UNITEOF
  systemctl daemon-reload
  systemctl enable cast-imaging-firewall.service >/dev/null 2>&1
  systemctl start cast-imaging-firewall.service
  echo "${c_grn}Persisted: systemd unit cast-imaging-firewall.service (enabled)${c_off}"
else
  warn "not persisted (-N): rules are lost at reboot."
fi

echo; echo "Verify from a non-Imaging machine: nc -zv <this-host> <port>   (8090: open, others: timeout)"
echo "Drop counter / rules:              sudo $0 status"
echo "Revert:                            sudo $0 remove"