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

# ThreeGate DOCKER-USER egress enforcement (clean)
#
# Block outbound internet egress from ThreeGate internal container networks.
# Allow ONLY the proxy (or egressnet subnet) to reach DNS + HTTPS.
#
# Recommended: pin explicit IPAM subnets and PROXY_IP in docker-compose.

CHAIN="DOCKER-USER"

# Operator settings (override via environment)
LLMNET_SUBNET="${LLMNET_SUBNET:-172.18.0.0/16}"
FETchnet_SUBNET="${FETchnet_SUBNET:-172.19.0.0/16}"
EGRESSNET_SUBNET="${EGRESSNET_SUBNET:-172.20.0.0/16}"

PROXY_IP="${PROXY_IP:-}"      # best: pin via IPAM
DNS_1="${DNS_1:-1.1.1.1}"
DNS_2="${DNS_2:-8.8.8.8}"

need_root() {
  if [[ "${EUID}" -ne 0 ]]; then
    echo "ERROR: must run as root" >&2
    exit 1
  fi
}

ensure_chain() {
  iptables -nL "${CHAIN}" >/dev/null 2>&1 || iptables -N "${CHAIN}"
  if ! iptables -C "${CHAIN}" -j RETURN >/dev/null 2>&1; then
    iptables -A "${CHAIN}" -j RETURN
  fi
}

reset_chain() {
  iptables -F "${CHAIN}"
  iptables -A "${CHAIN}" -j RETURN
}

insert_before_return() {
  local last
  last="$(iptables -nL "${CHAIN}" --line-numbers | tail -n 1 | awk '{print $1}')"
  iptables -I "${CHAIN}" "${last}" "$@"
}

main() {
  need_root
  ensure_chain
  reset_chain

  # Allow established traffic
  insert_before_return -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT

  # Allow proxy egress to HTTPS + DNS
  if [[ -n "${PROXY_IP}" ]]; then
    insert_before_return -s "${PROXY_IP}" -p tcp --dport 443 -j ACCEPT
    insert_before_return -s "${PROXY_IP}" -p udp -d "${DNS_1}" --dport 53 -j ACCEPT
    insert_before_return -s "${PROXY_IP}" -p udp -d "${DNS_2}" --dport 53 -j ACCEPT
    insert_before_return -s "${PROXY_IP}" -p tcp -d "${DNS_1}" --dport 53 -j ACCEPT
    insert_before_return -s "${PROXY_IP}" -p tcp -d "${DNS_2}" --dport 53 -j ACCEPT
  else
    echo "WARN: PROXY_IP not set. Allowing egress for sources in EGRESSNET_SUBNET=${EGRESSNET_SUBNET}." >&2
    insert_before_return -s "${EGRESSNET_SUBNET}" -p tcp --dport 443 -j ACCEPT
    insert_before_return -s "${EGRESSNET_SUBNET}" -p udp --dport 53 -j ACCEPT
    insert_before_return -s "${EGRESSNET_SUBNET}" -p tcp --dport 53 -j ACCEPT
  fi

  # Default-deny NEW outbound connections from internal networks
  insert_before_return -s "${LLMNET_SUBNET}" -m conntrack --ctstate NEW -j REJECT
  insert_before_return -s "${FETchnet_SUBNET}" -m conntrack --ctstate NEW -j REJECT

  echo "Applied ThreeGate DOCKER-USER egress policy."
  echo "  LLMNET_SUBNET=${LLMNET_SUBNET}"
  echo "  FETchnet_SUBNET=${FETchnet_SUBNET}"
  echo "  EGRESSNET_SUBNET=${EGRESSNET_SUBNET}"
  echo "  PROXY_IP=${PROXY_IP:-<unset>}"
  echo "  DNS_1=${DNS_1} DNS_2=${DNS_2}"
}

main "$@"