#!/bin/bash

# DartNode SMTP Gate Installer & Manager
# One-liner: curl -sSL https://pkg.dev.snaju.com/dn/smtp-gate.sh | bash -s install

PKG_BASE="https://pkg.dev.snaju.com/dn"

# Paths
SMTP_GATE_BIN="/usr/local/bin/smtp-gate"
SMTP_GATE_CONFIG_DIR="/etc/smtp-gate"
SMTP_GATE_SERVICE="smtp-gate"
SCRIPT_INSTALL_PATH="/usr/local/bin/smtp-gate-ctl"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
DIM='\033[2m'
NC='\033[0m'

# ============================================
# Utility Functions
# ============================================

function check_root() {
  if [ "$EUID" -ne 0 ]; then
    echo -e "${RED}Error: This command requires root privileges.${NC}"
    echo "Please run with sudo or as root."
    exit 1
  fi
}

function check_linux() {
  if [ "$(uname -s)" != "Linux" ]; then
    echo -e "${RED}Error: SMTP Gate is only supported on Linux.${NC}"
    exit 1
  fi
}

function print_step() {
  echo -e "${BLUE}=> ${NC}${BOLD}$1${NC}"
}

function print_success() {
  echo -e "${GREEN}✓  $1${NC}"
}

function print_error() {
  echo -e "${RED}✗  $1${NC}"
}

function print_warning() {
  echo -e "${YELLOW}!  $1${NC}"
}

function print_info() {
  echo -e "${DIM}   $1${NC}"
}

function header() {
  echo ""
  echo -e "${BOLD}DartNode SMTP Gate${NC}"
  echo -e "${DIM}Outbound SMTP Policy Enforcement (nfqueue)${NC}"
  echo ""
}

# ============================================
# Network Setup (nftables + nfqueue + IP forwarding)
# ============================================

function do_setup_network() {
  check_root
  check_linux

  print_step "Setting up nftables and IP forwarding for nfqueue filter mode..."

  # Load nfnetlink_queue kernel module
  print_step "Loading nfnetlink_queue kernel module..."
  modprobe nfnetlink_queue
  if [ $? -eq 0 ]; then
    print_success "nfnetlink_queue module loaded."
  else
    print_error "Failed to load nfnetlink_queue module."
    exit 1
  fi

  # Persist kernel module
  if ! grep -q 'nfnetlink_queue' /etc/modules-load.d/*.conf 2>/dev/null; then
    echo "nfnetlink_queue" >> /etc/modules-load.d/smtp-gate.conf
    print_success "nfnetlink_queue module set to load at boot."
  fi

  # Enable IP forwarding
  print_step "Enabling IP forwarding..."
  sysctl -w net.ipv4.ip_forward=1 >/dev/null
  sysctl -w net.ipv6.conf.all.forwarding=1 >/dev/null

  # Persist sysctl settings
  cat > /etc/sysctl.d/99-smtp-gate.conf <<-'EOF'
net.ipv4.ip_forward=1
net.ipv6.conf.all.forwarding=1
net.core.rmem_max=8388608
net.core.rmem_default=4194304
EOF
  sysctl --system >/dev/null 2>&1
  print_success "IP forwarding enabled and persisted."

  # Write nftables rules
  print_step "Writing nftables rules..."
  mkdir -p ${SMTP_GATE_CONFIG_DIR}
  cat > ${SMTP_GATE_CONFIG_DIR}/nftables.conf <<-'NFT'
#!/usr/sbin/nft -f

table inet smtp_gate {
	set blocked_customers {
		type ipv4_addr
	}

	set ratelimited_customers {
		type ipv4_addr
	}

	set blocked_customers_v6 {
		type ipv6_addr
	}

	set ratelimited_customers_v6 {
		type ipv6_addr
	}

	chain prerouting {
		type filter hook prerouting priority mangle; policy accept;

		# Skip packets already processed (marked DSCP AF11)
		ip dscp af11 accept
		ip6 dscp af11 accept
	}

	chain forward {
		type filter hook forward priority filter; policy accept;

		# Fast-path block: drop all SMTP from blocked customer IPs
		ip saddr @blocked_customers tcp dport 25 drop
		ip6 saddr @blocked_customers_v6 tcp dport 25 drop

		# Fast-path rate limit: reject SMTP from rate-limited customer IPs
		ip saddr @ratelimited_customers tcp dport 25 reject with tcp reset
		ip6 saddr @ratelimited_customers_v6 tcp dport 25 reject with tcp reset

		# Send all other port 25 traffic to nfqueue for inspection
		tcp dport 25 queue num 0 bypass
	}

	chain postrouting {
		type filter hook postrouting priority mangle; policy accept;

		# Mark forwarded SMTP with DSCP AF11 so edge router doesn't re-route
		tcp dport 25 ip dscp set af11
		tcp dport 25 ip6 dscp set af11
	}
}
NFT

  # Apply nftables rules (flush smtp_gate table and reload)
  nft delete table inet smtp_gate 2>/dev/null
  nft -f ${SMTP_GATE_CONFIG_DIR}/nftables.conf
  if [ $? -eq 0 ]; then
    print_success "nftables rules applied."
  else
    print_error "Failed to apply nftables rules."
    exit 1
  fi

  # Install systemd service for network setup persistence across reboots
  print_step "Installing network setup persistence service..."
  cat > /etc/systemd/system/smtp-gate-network.service <<-EOF
[Unit]
Description=SMTP Gate network setup (nfqueue + IP forwarding)
Before=smtp-gate.service
After=network-online.target

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/bin/sh -c 'modprobe nfnetlink_queue; sysctl -w net.ipv4.ip_forward=1 >/dev/null; sysctl -w net.ipv6.conf.all.forwarding=1 >/dev/null; nft -f ${SMTP_GATE_CONFIG_DIR}/nftables.conf'
ExecStop=/bin/sh -c 'nft delete table inet smtp_gate 2>/dev/null || true'

[Install]
WantedBy=multi-user.target
EOF

  # Clean up old TPROXY service if it exists
  if [ -f /etc/systemd/system/smtp-tproxy-routes.service ]; then
    print_info "Removing old TPROXY routing service..."
    systemctl stop smtp-tproxy-routes 2>/dev/null
    systemctl disable smtp-tproxy-routes 2>/dev/null
    rm -f /etc/systemd/system/smtp-tproxy-routes.service
    # Clean up old TPROXY routing rules
    ip rule del fwmark 1 lookup 100 2>/dev/null
    ip -6 rule del fwmark 1 lookup 100 2>/dev/null
  fi

  systemctl daemon-reload
  systemctl enable smtp-gate-network
  systemctl start smtp-gate-network

  print_success "Network setup complete (nfqueue + IP forwarding)."
  echo ""
}

# ============================================
# Install
# ============================================

function do_install() {
  check_root
  check_linux

  print_step "Installing SMTP Gate..."
  echo ""

  # Create directories
  mkdir -p ${SMTP_GATE_CONFIG_DIR}

  # Stop existing service if running
  if systemctl is-active --quiet ${SMTP_GATE_SERVICE} 2>/dev/null; then
    print_info "Stopping existing smtp-gate service..."
    systemctl stop ${SMTP_GATE_SERVICE}
  fi

  # Create smtpgate system user if it doesn't exist
  if ! id smtpgate &>/dev/null; then
    print_step "Creating smtpgate system user..."
    useradd -r -s /usr/sbin/nologin -d /nonexistent smtpgate
    print_success "smtpgate user created."
  fi

  # Download binary
  print_step "Downloading smtp-gate binary..."
  wget -q --show-progress -O ${SMTP_GATE_BIN} ${PKG_BASE}/smtp-gate
  if [ $? -ne 0 ]; then
    print_error "Failed to download smtp-gate binary."
    exit 1
  fi
  chmod +x ${SMTP_GATE_BIN}
  chown smtpgate:smtpgate ${SMTP_GATE_BIN}

  # Write default config if none exists
  if [ ! -f "${SMTP_GATE_CONFIG_DIR}/config.yaml" ]; then
    print_step "Writing default config..."
    cat > "${SMTP_GATE_CONFIG_DIR}/config.yaml" <<-'YAML'
mode: "filter"
nfqueue_num: 0
listen: ":2525"
api_listen: ":9090"

redis:
  addr: "localhost:6379"
  password: ""
  db: 0
  pool_size: 20

default_policy: "block"

# Per-IP hourly limit (prevents a single VM from spamming)
default_ip_limit: 50
default_ip_window_secs: 3600

# Per-customer hourly limit (aggregate across all IPs)
default_limit: 100
default_window_secs: 3600

# Per-customer daily limit (prevents sustained abuse)
default_daily_limit: 500

max_connections: 5000
shutdown_timeout_secs: 30
log_level: "info"

nftables:
  table: "inet smtp_gate"
  blocked_set: "blocked_customers"
  ratelimited_set: "ratelimited_customers"
  blocked_set_v6: "blocked_customers_v6"
  ratelimited_set_v6: "ratelimited_customers_v6"
YAML
    print_info "Edit ${SMTP_GATE_CONFIG_DIR}/config.yaml to configure Redis and policy settings."
  else
    print_info "Existing config preserved at ${SMTP_GATE_CONFIG_DIR}/config.yaml"
  fi

  # Ensure config directory is readable by smtpgate user
  chown -R smtpgate:smtpgate ${SMTP_GATE_CONFIG_DIR}

  # Install systemd service
  print_step "Installing systemd service..."
  cat > "/etc/systemd/system/${SMTP_GATE_SERVICE}.service" <<-EOF
[Unit]
Description=DartNode SMTP Gate — Outbound SMTP Policy Enforcement (nfqueue)
After=network-online.target redis.service smtp-gate-network.service
Wants=network-online.target
Requires=smtp-gate-network.service

[Service]
Type=simple
User=smtpgate
Group=smtpgate
ExecStart=${SMTP_GATE_BIN} -config ${SMTP_GATE_CONFIG_DIR}/config.yaml
Restart=on-failure
RestartSec=5

AmbientCapabilities=CAP_NET_ADMIN CAP_NET_RAW
CapabilityBoundingSet=CAP_NET_ADMIN CAP_NET_RAW
LimitNOFILE=65535

StandardOutput=journal
StandardError=journal
SyslogIdentifier=smtp-gate

[Install]
WantedBy=multi-user.target
EOF

  systemctl daemon-reload
  systemctl enable ${SMTP_GATE_SERVICE}

  # Install this script as smtp-gate-ctl
  print_step "Installing smtp-gate-ctl management command..."
  if [ -f "$0" ] && [ "$0" != "bash" ] && [ "$0" != "/dev/stdin" ]; then
    cp "$0" ${SCRIPT_INSTALL_PATH}
    chmod +x ${SCRIPT_INSTALL_PATH}
    print_success "Installed: ${SCRIPT_INSTALL_PATH}"
  else
    # Piped install — download the script
    wget -q -O ${SCRIPT_INSTALL_PATH} ${PKG_BASE}/smtp-gate.sh
    chmod +x ${SCRIPT_INSTALL_PATH}
    print_success "Installed: ${SCRIPT_INSTALL_PATH}"
  fi

  # Set up nftables and IP forwarding
  do_setup_network

  echo ""
  echo -e "${GREEN}============================================${NC}"
  echo -e "${GREEN}      SMTP Gate Installation Complete!      ${NC}"
  echo -e "${GREEN}============================================${NC}"
  echo ""
  echo "Config:   ${SMTP_GATE_CONFIG_DIR}/config.yaml"
  echo "nftables: ${SMTP_GATE_CONFIG_DIR}/nftables.conf"
  echo "Binary:   ${SMTP_GATE_BIN}"
  echo "Service:  ${SMTP_GATE_SERVICE}"
  echo "Manager:  ${SCRIPT_INSTALL_PATH}"
  echo ""
  echo -e "${BOLD}Architecture:${NC}"
  echo "  Mode: nfqueue packet filter (transparent forwarding)"
  echo "  Customer IPs are preserved end-to-end"
  echo "  Only MAIL FROM commands are counted for rate limiting"
  echo ""
  echo -e "${BOLD}Rate Limit Defaults:${NC}"
  echo "  Per-IP hourly:      50 emails/hour  (catches single-VM spam)"
  echo "  Per-customer hourly: 100 emails/hour (aggregate across all VMs)"
  echo "  Per-customer daily:  500 emails/day  (prevents sustained abuse)"
  echo ""
  echo "Commands:"
  echo "  smtp-gate-ctl start         - Start smtp-gate"
  echo "  smtp-gate-ctl stop          - Stop smtp-gate"
  echo "  smtp-gate-ctl restart       - Restart smtp-gate"
  echo "  smtp-gate-ctl status        - Check status & health"
  echo "  smtp-gate-ctl log           - Follow logs"
  echo "  smtp-gate-ctl update        - Update binary"
  echo "  smtp-gate-ctl config        - Edit config"
  echo "  smtp-gate-ctl setup-network - Set up nftables + nfqueue"
  echo "  smtp-gate-ctl uninstall     - Remove smtp-gate"
  echo ""

  # Ask to start
  if [ -t 0 ]; then
    read -p "Start smtp-gate now? (y/n) [y]: " start_now
    start_now=${start_now:-y}
    if [[ $start_now =~ ^[Yy]$ ]]; then
      do_start
    fi
  fi
}

# ============================================
# Service Management
# ============================================

function do_start() {
  check_root
  if [ ! -x "${SMTP_GATE_BIN}" ]; then
    print_error "SMTP Gate not installed. Run: smtp-gate-ctl install"
    exit 1
  fi
  print_step "Starting smtp-gate..."
  systemctl start ${SMTP_GATE_SERVICE}
  sleep 1
  if systemctl is-active --quiet ${SMTP_GATE_SERVICE}; then
    print_success "smtp-gate is running."
  else
    print_error "smtp-gate failed to start. Check: journalctl -u ${SMTP_GATE_SERVICE} -n 20"
  fi
}

function do_stop() {
  check_root
  print_step "Stopping smtp-gate..."
  systemctl stop ${SMTP_GATE_SERVICE}
  print_success "smtp-gate stopped."
}

function do_restart() {
  check_root
  print_step "Restarting smtp-gate..."
  systemctl restart ${SMTP_GATE_SERVICE}
  sleep 1
  if systemctl is-active --quiet ${SMTP_GATE_SERVICE}; then
    print_success "smtp-gate restarted."
  else
    print_error "smtp-gate failed to start. Check: journalctl -u ${SMTP_GATE_SERVICE} -n 20"
  fi
}

function do_status() {
  if [ ! -x "${SMTP_GATE_BIN}" ]; then
    print_error "SMTP Gate not installed."
    exit 1
  fi

  echo -e "${BOLD}SMTP Gate Status${NC}"
  echo "---"

  echo -n "Service:  "
  if systemctl is-active --quiet ${SMTP_GATE_SERVICE} 2>/dev/null; then
    echo -e "${GREEN}running${NC}"
  else
    echo -e "${RED}stopped${NC}"
  fi

  echo -n "Enabled:  "
  if systemctl is-enabled --quiet ${SMTP_GATE_SERVICE} 2>/dev/null; then
    echo -e "${GREEN}yes${NC}"
  else
    echo -e "${YELLOW}no${NC}"
  fi

  echo -n "Network:  "
  if systemctl is-active --quiet smtp-gate-network 2>/dev/null; then
    echo -e "${GREEN}ready${NC}"
  else
    echo -e "${RED}not configured${NC}"
  fi

  echo "Binary:   ${SMTP_GATE_BIN}"
  echo "Config:   ${SMTP_GATE_CONFIG_DIR}/config.yaml"

  if [ -x "${SMTP_GATE_BIN}" ]; then
    echo -n "Version:  "
    ${SMTP_GATE_BIN} -version 2>/dev/null || echo "unknown"
  fi

  # Mode
  echo -n "Mode:     "
  if [ -f "${SMTP_GATE_CONFIG_DIR}/config.yaml" ]; then
    MODE=$(grep -oP 'mode:\s*"\K[^"]+' ${SMTP_GATE_CONFIG_DIR}/config.yaml 2>/dev/null || echo "filter")
    echo "${MODE}"
  else
    echo "unknown"
  fi

  # Health check
  echo ""
  echo -n "Health:   "
  HEALTH=$(curl -s --connect-timeout 2 http://localhost:9090/health 2>/dev/null)
  if [ -n "$HEALTH" ]; then
    echo -e "${GREEN}${HEALTH}${NC}"
  else
    echo -e "${RED}unreachable${NC}"
  fi

  # IP forwarding check
  echo ""
  echo -e "${BOLD}Network:${NC}"
  echo -n "  IPv4 forwarding: "
  FWD=$(cat /proc/sys/net/ipv4/ip_forward 2>/dev/null)
  if [ "$FWD" = "1" ]; then
    echo -e "${GREEN}enabled${NC}"
  else
    echo -e "${RED}disabled${NC}"
  fi

  echo -n "  nfnetlink_queue: "
  if lsmod | grep -q nfnetlink_queue 2>/dev/null; then
    echo -e "${GREEN}loaded${NC}"
  else
    echo -e "${RED}not loaded${NC}"
  fi

  echo -n "  nftables table:  "
  if nft list table inet smtp_gate >/dev/null 2>&1; then
    echo -e "${GREEN}present${NC}"
  else
    echo -e "${RED}missing${NC}"
  fi

  # Show rate limit defaults from config
  echo ""
  echo -e "${BOLD}Rate Limits (from config):${NC}"
  if [ -f "${SMTP_GATE_CONFIG_DIR}/config.yaml" ]; then
    IP_LIMIT=$(grep -oP 'default_ip_limit:\s*\K\d+' ${SMTP_GATE_CONFIG_DIR}/config.yaml 2>/dev/null || echo "50")
    HOURLY_LIMIT=$(grep -oP 'default_limit:\s*\K\d+' ${SMTP_GATE_CONFIG_DIR}/config.yaml 2>/dev/null || echo "100")
    DAILY_LIMIT=$(grep -oP 'default_daily_limit:\s*\K\d+' ${SMTP_GATE_CONFIG_DIR}/config.yaml 2>/dev/null || echo "500")
    echo "  Per-IP hourly:       ${IP_LIMIT} emails/hour"
    echo "  Per-customer hourly: ${HOURLY_LIMIT} emails/hour"
    echo "  Per-customer daily:  ${DAILY_LIMIT} emails/day"
  else
    echo "  (config not found)"
  fi
}

function do_log() {
  journalctl -u ${SMTP_GATE_SERVICE} -f
}

function do_update() {
  check_root
  print_step "Updating smtp-gate..."

  WAS_RUNNING=false
  if systemctl is-active --quiet ${SMTP_GATE_SERVICE} 2>/dev/null; then
    WAS_RUNNING=true
    systemctl stop ${SMTP_GATE_SERVICE}
  fi

  wget -q --show-progress -O ${SMTP_GATE_BIN} ${PKG_BASE}/smtp-gate
  if [ $? -ne 0 ]; then
    print_error "Failed to download update."
    if [ "$WAS_RUNNING" = true ]; then
      systemctl start ${SMTP_GATE_SERVICE}
    fi
    exit 1
  fi
  chmod +x ${SMTP_GATE_BIN}
  chown smtpgate:smtpgate ${SMTP_GATE_BIN}

  # Also update this management script
  wget -q -O ${SCRIPT_INSTALL_PATH} ${PKG_BASE}/smtp-gate.sh 2>/dev/null
  chmod +x ${SCRIPT_INSTALL_PATH} 2>/dev/null

  if [ "$WAS_RUNNING" = true ]; then
    systemctl start ${SMTP_GATE_SERVICE}
  fi

  print_success "smtp-gate updated."
}

function do_config() {
  if [ ! -f "${SMTP_GATE_CONFIG_DIR}/config.yaml" ]; then
    print_error "Config not found at ${SMTP_GATE_CONFIG_DIR}/config.yaml"
    exit 1
  fi

  # Prefer $EDITOR, fall back to nano, then vi
  EDITOR_CMD="${EDITOR:-nano}"
  if ! command -v "$EDITOR_CMD" &>/dev/null; then
    EDITOR_CMD="vi"
  fi

  $EDITOR_CMD "${SMTP_GATE_CONFIG_DIR}/config.yaml"

  echo ""
  print_info "Restart smtp-gate to apply changes: smtp-gate-ctl restart"
}

function do_uninstall() {
  check_root

  echo -e "${RED}This will completely remove SMTP Gate from this server.${NC}"
  if [ -t 0 ]; then
    read -p "Are you sure? (y/n) " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
      echo "Cancelled."
      exit 0
    fi
  fi

  print_step "Uninstalling smtp-gate..."

  systemctl stop ${SMTP_GATE_SERVICE} 2>/dev/null
  systemctl disable ${SMTP_GATE_SERVICE} 2>/dev/null
  rm -f "/etc/systemd/system/${SMTP_GATE_SERVICE}.service"

  # Clean up network service
  systemctl stop smtp-gate-network 2>/dev/null
  systemctl disable smtp-gate-network 2>/dev/null
  rm -f /etc/systemd/system/smtp-gate-network.service

  # Clean up old TPROXY service if still present
  systemctl stop smtp-tproxy-routes 2>/dev/null
  systemctl disable smtp-tproxy-routes 2>/dev/null
  rm -f /etc/systemd/system/smtp-tproxy-routes.service

  # Remove nftables table
  nft delete table inet smtp_gate 2>/dev/null

  # Clean up old TPROXY routing rules
  ip rule del fwmark 1 lookup 100 2>/dev/null
  ip -6 rule del fwmark 1 lookup 100 2>/dev/null

  systemctl daemon-reload
  rm -f ${SMTP_GATE_BIN}
  rm -f ${SCRIPT_INSTALL_PATH}
  rm -f /etc/sysctl.d/99-smtp-gate.conf
  rm -f /etc/modules-load.d/smtp-gate.conf

  echo ""
  print_info "Config preserved at ${SMTP_GATE_CONFIG_DIR}/"
  print_info "To remove config: rm -rf ${SMTP_GATE_CONFIG_DIR}"
  print_success "smtp-gate uninstalled."
}

# ============================================
# Help
# ============================================

function help() {
  echo -e "${BOLD}Usage:${NC} smtp-gate-ctl <command>"
  echo ""
  echo -e "${BOLD}Installation:${NC}"
  echo "  install     - Install SMTP Gate daemon and service"
  echo ""
  echo -e "${BOLD}Service Management:${NC}"
  echo "  start       - Start smtp-gate"
  echo "  stop        - Stop smtp-gate"
  echo "  restart     - Restart smtp-gate"
  echo "  status      - Show status, health, and rate limits"
  echo "  log         - Follow smtp-gate logs"
  echo ""
  echo -e "${BOLD}Configuration:${NC}"
  echo "  config        - Edit config file"
  echo "  update        - Update smtp-gate binary"
  echo "  setup-network - Set up nftables + nfqueue + IP forwarding"
  echo "  uninstall     - Remove smtp-gate"
  echo ""
  echo -e "${BOLD}Architecture:${NC}"
  echo "  nfqueue packet filter — transparent SMTP forwarding"
  echo "  Customer source IPs preserved end-to-end"
  echo "  Only MAIL FROM commands counted for rate limiting"
  echo ""
  echo -e "${BOLD}Rate Limit Defaults:${NC}"
  echo "  Per-IP hourly:       50 emails/hour"
  echo "  Per-customer hourly: 100 emails/hour"
  echo "  Per-customer daily:  500 emails/day"
  echo ""
  echo -e "${BOLD}One-liner install:${NC}"
  echo "  curl -sSL ${PKG_BASE}/smtp-gate.sh | bash -s install"
  echo ""
}

# ============================================
# Main
# ============================================

header

COMMAND="${1:-help}"

case ${COMMAND} in
  install)
    do_install
    ;;
  start)
    do_start
    ;;
  stop)
    do_stop
    ;;
  restart)
    do_restart
    ;;
  status)
    do_status
    ;;
  log|logs)
    do_log
    ;;
  update)
    do_update
    ;;
  config)
    do_config
    ;;
  setup-network)
    do_setup_network
    ;;
  uninstall|remove)
    do_uninstall
    ;;
  help|--help|-h)
    help
    ;;
  *)
    echo "Unknown command: ${COMMAND}"
    echo ""
    help
    exit 1
    ;;
esac