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

usage() {
  cat <<'EOH'
Usage: start-transparent.sh [-v|--verbose] [--enable-udp] [--capture-uid UID] [--kernel auto|20|24] [config_path]

Options:
  -v, --verbose        启动后实时输出 mynetspeeder 日志
  --capture-uid UID    指定时只接管该 UID；不指定时接管所有用户流量
  --enable-udp         额外启用 UDP 透明接管（实验性，默认关闭）
  --kernel MODE        指定内核优化模式：auto|20|24，默认 auto
  -h, --help           显示帮助
EOH
}

VERBOSE=0
ENABLE_UDP=0
KERNEL_MODE="${MYNETSPEEDER_KERNEL_MODE:-auto}"
CONFIG_PATH="/home/mynetspeeder/config.json"
CAPTURE_UID="${MYNETSPEEDER_CAPTURE_UID:-}"
UDP_CAPTURE_EFFECTIVE=0

while [[ $# -gt 0 ]]; do
  case "$1" in
    -v|--verbose)
      VERBOSE=1; shift ;;
    --enable-udp)
      ENABLE_UDP=1; shift ;;
    --capture-uid)
      CAPTURE_UID="${2:-}"
      [[ -n "$CAPTURE_UID" ]] || { echo "missing value for --capture-uid"; exit 1; }
      shift 2 ;;
    --kernel)
      KERNEL_MODE="${2:-}"
      [[ -n "$KERNEL_MODE" ]] || { echo "missing value for --kernel"; exit 1; }
      shift 2 ;;
    -h|--help)
      usage; exit 0 ;;
    *)
      CONFIG_PATH="$1"; shift ;;
  esac
done

LISTEN_HOST="${MYNETSPEEDER_LISTEN_HOST:-127.0.0.1}"
LISTEN_PORT="${MYNETSPEEDER_LISTEN_PORT:-19080}"
RUNTIME_USER="${MYNETSPEEDER_USER:-mynetspeeder}"
PID_FILE="/var/run/mynetspeeder-edge.pid"
SOCKS_PID_FILE="/var/run/mynetspeeder-socks.pid"
LOG_FILE="/var/log/mynetspeeder-edge.log"
SOCKS_LOG_FILE="/var/log/mynetspeeder-socks.log"
LOG_MAX_MB="${MYNETSPEEDER_LOG_MAX_MB:-50}"
LOG_BACKUPS="${MYNETSPEEDER_LOG_BACKUPS:-3}"
CHAIN4="MYNETSPEEDER"
CHAIN6="MYNETSPEEDER6"
SSH_PORTS="${MYNETSPEEDER_SSH_PORTS:-}"
SELF_EXCLUDE_V4="127.0.0.0/8 169.254.0.0/16"
SELF_EXCLUDE_V6="::1/128 fe80::/10"

if [[ $EUID -ne 0 ]]; then echo "need root"; exit 1; fi
if [[ ! -f "$CONFIG_PATH" ]]; then echo "config not found: $CONFIG_PATH"; exit 1; fi
if [[ -n "$CAPTURE_UID" ]] && ! [[ "$CAPTURE_UID" =~ ^[0-9]+$ ]]; then echo "capture uid must be numeric"; exit 1; fi
case "$KERNEL_MODE" in
  auto|20|24) ;;
  *) echo "invalid kernel mode: $KERNEL_MODE"; exit 1 ;;
esac

if [[ "$KERNEL_MODE" == "auto" ]]; then
  if [[ -f /etc/os-release ]] && grep -q '^VERSION_ID="24' /etc/os-release; then
    KERNEL_MODE="24"
  else
    KERNEL_MODE="20"
  fi
fi

if [[ -z "$SSH_PORTS" && -n "${SSH_CONNECTION:-}" ]]; then
  SSH_PORTS="${SSH_CONNECTION##* }"
fi

SSH_PORT_ARRAY=()
if [[ -n "$SSH_PORTS" ]]; then
  IFS=',' read -r -a SSH_PORT_ARRAY <<< "$SSH_PORTS"
  for ssh_port in "${SSH_PORT_ARRAY[@]}"; do
    [[ "$ssh_port" =~ ^[0-9]+$ ]] || { echo "ssh ports must be numeric, got: $ssh_port"; exit 1; }
  done
fi

id -u "$RUNTIME_USER" >/dev/null 2>&1 || useradd --system --no-create-home --shell /usr/sbin/nologin "$RUNTIME_USER"
mkdir -p /var/log
touch "$LOG_FILE"
chown "$RUNTIME_USER":"$RUNTIME_USER" "$LOG_FILE"
touch "$SOCKS_LOG_FILE"
chown "$RUNTIME_USER":"$RUNTIME_USER" "$SOCKS_LOG_FILE"
if ! [[ "$LOG_MAX_MB" =~ ^[0-9]+$ ]] || ! [[ "$LOG_BACKUPS" =~ ^[0-9]+$ ]]; then echo "log limits must be numeric"; exit 1; fi
LOG_MAX_BYTES=$((LOG_MAX_MB * 1024 * 1024))

IPTABLES_BACKEND="unknown"
if iptables --version 2>/dev/null | grep -qi 'nf_tables'; then
  IPTABLES_BACKEND="nf_tables"
else
  IPTABLES_BACKEND="legacy"
fi

ensure_rule() {
  local cmd="$1"
  local table="$2"
  local chain="$3"
  shift 3
  if ! "$cmd" -t "$table" -C "$chain" "$@" >/dev/null 2>&1; then
    "$cmd" -t "$table" -A "$chain" "$@"
  fi
}

add_exclusions_v4() {
  iptables -t nat -A "$CHAIN4" -m addrtype --dst-type LOCAL -j RETURN
  for cidr in $SELF_EXCLUDE_V4; do
    iptables -t nat -A "$CHAIN4" -d "$cidr" -j RETURN
  done
  iptables -t nat -A "$CHAIN4" -m owner --uid-owner "$RUNTIME_USER" -j RETURN
  for ssh_port in "${SSH_PORT_ARRAY[@]}"; do
    iptables -t nat -A "$CHAIN4" -p tcp --sport "$ssh_port" -j RETURN
  done
  while read -r host; do
    [[ -n "$host" && "$host" != *:* ]] && iptables -t nat -A "$CHAIN4" -d "$host" -j RETURN
  done < <(python3 - <<'PY' "$CONFIG_PATH"
import json, sys
cfg = json.load(open(sys.argv[1]))
for relay in cfg.get('relays', []):
    print(relay['host'])
PY
)
}

add_exclusions_v6() {
  ip6tables -t nat -A "$CHAIN6" -m addrtype --dst-type LOCAL -j RETURN
  for cidr in $SELF_EXCLUDE_V6; do
    ip6tables -t nat -A "$CHAIN6" -d "$cidr" -j RETURN
  done
  ip6tables -t nat -A "$CHAIN6" -m owner --uid-owner "$RUNTIME_USER" -j RETURN
  for ssh_port in "${SSH_PORT_ARRAY[@]}"; do
    ip6tables -t nat -A "$CHAIN6" -p tcp --sport "$ssh_port" -j RETURN
  done
  while read -r host; do
    [[ -n "$host" && "$host" == *:* ]] && ip6tables -t nat -A "$CHAIN6" -d "$host" -j RETURN
  done < <(python3 - <<'PY' "$CONFIG_PATH"
import json, sys
cfg = json.load(open(sys.argv[1]))
for relay in cfg.get('relays', []):
    print(relay['host'])
PY
)
}

add_udp_exclusions_v4() {
  iptables -t nat -A "$CHAIN4" -p udp -m conntrack --ctstate ESTABLISHED,RELATED -j RETURN
}

add_udp_exclusions_v6() {
  ip6tables -t nat -A "$CHAIN6" -p udp -m conntrack --ctstate ESTABLISHED,RELATED -j RETURN
}

pkill -f 'python3 -m mynetspeeder edge' || true
pkill -f 'python3 -m mynetspeeder socks' || true

SOCKS_HOST=$(python3 - <<'PY' "$CONFIG_PATH"
import json, sys
cfg = json.load(open(sys.argv[1]))
print(cfg.get('socks_host', '127.0.0.1'))
PY
)
SOCKS_PORT=$(python3 - <<'PY' "$CONFIG_PATH"
import json, sys
cfg = json.load(open(sys.argv[1]))
print(int(cfg.get('socks_port', 0) or 0))
PY
)

if [[ "$SOCKS_PORT" -gt 0 ]]; then
  runuser -u "$RUNTIME_USER" -- bash -lc "export PYTHONUNBUFFERED=1; export PYTHONPATH=/home; cd /home && exec nohup python3 -m mynetspeeder socks --listen-host ${SOCKS_HOST} --listen-port ${SOCKS_PORT} --config ${CONFIG_PATH} 2>&1 | python3 /home/mynetspeeder/scripts/rotate-log.py ${SOCKS_LOG_FILE} ${LOG_MAX_BYTES} ${LOG_BACKUPS}" &
  SOCKS_PID=$!
  echo "$SOCKS_PID" > "$SOCKS_PID_FILE"
  sleep 1
  ss -lntp | grep -qE "${SOCKS_HOST//./\\.}:${SOCKS_PORT}( |$)" || { echo "socks failed to listen"; tail -n 50 "$SOCKS_LOG_FILE" || true; exit 1; }
fi

if [[ "$ENABLE_UDP" == "1" && "$SOCKS_PORT" -gt 0 ]]; then
  echo "udp transparent capture disabled: socks5 is enabled, UDP will use socks only"
else
  UDP_CAPTURE_EFFECTIVE="$ENABLE_UDP"
fi

EDGE_UDP_FLAG=""
if [[ "$UDP_CAPTURE_EFFECTIVE" == "1" ]]; then
  EDGE_UDP_FLAG="--enable-udp"
fi
runuser -u "$RUNTIME_USER" -- bash -lc "export PYTHONUNBUFFERED=1; export PYTHONPATH=/home; cd /home && exec nohup python3 -m mynetspeeder edge --listen-host ${LISTEN_HOST} --listen-port ${LISTEN_PORT} --kernel ${KERNEL_MODE} --config ${CONFIG_PATH} ${EDGE_UDP_FLAG} 2>&1 | python3 /home/mynetspeeder/scripts/rotate-log.py ${LOG_FILE} ${LOG_MAX_BYTES} ${LOG_BACKUPS}" &
EDGE_PID=$!
echo "$EDGE_PID" > "$PID_FILE"
sleep 1
ss -ln | grep -qE "[:.]${LISTEN_PORT}( |$)" || { echo "edge failed to listen"; tail -n 50 "$LOG_FILE" || true; exit 1; }

iptables -t nat -N "$CHAIN4" 2>/dev/null || true
iptables -t nat -F "$CHAIN4"
add_exclusions_v4
if [[ -n "$CAPTURE_UID" ]]; then
  iptables -t nat -A "$CHAIN4" -p tcp -m owner --uid-owner "$CAPTURE_UID" -j REDIRECT --to-ports "$LISTEN_PORT"
else
  iptables -t nat -A "$CHAIN4" -p tcp -j REDIRECT --to-ports "$LISTEN_PORT"
fi
ensure_rule iptables nat OUTPUT -p tcp -j "$CHAIN4"
if [[ "$UDP_CAPTURE_EFFECTIVE" == "1" ]]; then
  add_udp_exclusions_v4
  if [[ -n "$CAPTURE_UID" ]]; then
    iptables -t nat -A "$CHAIN4" -p udp -m owner --uid-owner "$CAPTURE_UID" -j REDIRECT --to-ports "$LISTEN_PORT"
  else
    iptables -t nat -A "$CHAIN4" -p udp -j REDIRECT --to-ports "$LISTEN_PORT"
  fi
  ensure_rule iptables nat OUTPUT -p udp -j "$CHAIN4"
fi

IP6_ENABLED=0
IP6_NAT_SUPPORTED=0
if command -v ip6tables >/dev/null 2>&1; then
  if ip6tables -t nat -S >/dev/null 2>&1; then
    IP6_ENABLED=1
    IP6_NAT_SUPPORTED=1
    ip6tables -t nat -N "$CHAIN6" 2>/dev/null || true
    ip6tables -t nat -F "$CHAIN6"
    add_exclusions_v6
    if [[ -n "$CAPTURE_UID" ]]; then
      ip6tables -t nat -A "$CHAIN6" -p tcp -m owner --uid-owner "$CAPTURE_UID" -j REDIRECT --to-ports "$LISTEN_PORT"
    else
      ip6tables -t nat -A "$CHAIN6" -p tcp -j REDIRECT --to-ports "$LISTEN_PORT"
    fi
    ensure_rule ip6tables nat OUTPUT -p tcp -j "$CHAIN6"
    if [[ "$UDP_CAPTURE_EFFECTIVE" == "1" ]]; then
      add_udp_exclusions_v6
      if [[ -n "$CAPTURE_UID" ]]; then
        ip6tables -t nat -A "$CHAIN6" -p udp -m owner --uid-owner "$CAPTURE_UID" -j REDIRECT --to-ports "$LISTEN_PORT"
      else
        ip6tables -t nat -A "$CHAIN6" -p udp -j REDIRECT --to-ports "$LISTEN_PORT"
      fi
      ensure_rule ip6tables nat OUTPUT -p udp -j "$CHAIN6"
    fi
  else
    echo "ipv6 nat unavailable: ip6tables nat table not supported, skip ipv6 transparent rules"
  fi
fi

RULES_V4=$(iptables -t nat -S "$CHAIN4" 2>/dev/null | wc -l)
RULES_V6=0
if [[ "$IP6_ENABLED" == "1" && "$IP6_NAT_SUPPORTED" == "1" ]]; then
  RULES_V6=$(ip6tables -t nat -S "$CHAIN6" 2>/dev/null | wc -l)
fi
iptables -t nat -C OUTPUT -p tcp -j "$CHAIN4" >/dev/null 2>&1 || { echo "self-check failed: ipv4 tcp output hook missing"; exit 1; }
if [[ "$UDP_CAPTURE_EFFECTIVE" == "1" ]]; then
  iptables -t nat -C OUTPUT -p udp -j "$CHAIN4" >/dev/null 2>&1 || { echo "self-check failed: ipv4 udp output hook missing"; exit 1; }
fi
if [[ "$IP6_ENABLED" == "1" && "$IP6_NAT_SUPPORTED" == "1" ]]; then
  ip6tables -t nat -C OUTPUT -p tcp -j "$CHAIN6" >/dev/null 2>&1 || { echo "self-check failed: ipv6 tcp output hook missing"; exit 1; }
  if [[ "$UDP_CAPTURE_EFFECTIVE" == "1" ]]; then
    ip6tables -t nat -C OUTPUT -p udp -j "$CHAIN6" >/dev/null 2>&1 || { echo "self-check failed: ipv6 udp output hook missing"; exit 1; }
  fi
fi

echo "mynetspeeder transparent mode started on ${LISTEN_HOST}:${LISTEN_PORT}"
echo "kernel mode: $KERNEL_MODE"
echo "iptables backend: $IPTABLES_BACKEND"
if [[ -n "$CAPTURE_UID" ]]; then
  echo "capture uid: $CAPTURE_UID"
else
  echo "capture uid: all users"
fi
if [[ ${#SSH_PORT_ARRAY[@]} -gt 0 ]]; then
  echo "ssh exempt ports: ${SSH_PORT_ARRAY[*]}"
else
  echo "ssh exempt ports: none"
fi
echo "udp capture requested: $ENABLE_UDP"
echo "udp capture effective: $UDP_CAPTURE_EFFECTIVE"
if [[ "$SOCKS_PORT" -gt 0 ]]; then
  echo "socks5: ${SOCKS_HOST}:${SOCKS_PORT}"
  echo "socks log: $SOCKS_LOG_FILE"
else
  echo "socks5: disabled"
fi
echo "log file: $LOG_FILE"
echo "log max: ${LOG_MAX_MB}MB x ${LOG_BACKUPS}"
echo "ipv4 chain rules: $RULES_V4"
echo "ipv6 chain rules: $RULES_V6"

echo "self-check: ok"

if [[ "$VERBOSE" == "1" ]]; then
  echo "verbose mode: press Ctrl+C to stop viewing logs, service keeps running"
  exec tail -n 80 -f "$LOG_FILE"
fi