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

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

Options:
  -v, --verbose        启动后实时输出 mynetspeeder 日志
  --capture-uid UID    只接管该 UID 发起的 TCP 出站
  --enable-udp         额外启用 UDP 透明接管（实验性，默认关闭）
  -h, --help           显示帮助
EOF
}

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

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 ;;
    -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"
LOG_FILE="/var/log/mynetspeeder-edge.log"
CHAIN4="MYNETSPEEDER"
CHAIN6="MYNETSPEEDER6"

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 [[ -z "$CAPTURE_UID" ]]; then echo "refusing unsafe global capture"; exit 1; fi
if ! [[ "$CAPTURE_UID" =~ ^[0-9]+$ ]]; then echo "capture uid must be numeric"; exit 1; fi

id -u "$RUNTIME_USER" >/dev/null 2>&1 || useradd --system --no-create-home --shell /usr/sbin/nologin "$RUNTIME_USER"
mkdir -p /var/log
: > "$LOG_FILE"
chown "$RUNTIME_USER":"$RUNTIME_USER" "$LOG_FILE"

pkill -f 'python3 -m mynetspeeder edge' || true
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} --config ${CONFIG_PATH}" >>"$LOG_FILE" 2>&1 &
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"
iptables -t nat -A "$CHAIN4" -d 127.0.0.0/8 -j RETURN
iptables -t nat -A "$CHAIN4" -m owner --uid-owner "$RUNTIME_USER" -j RETURN
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
)
iptables -t nat -A "$CHAIN4" -p tcp -m owner --uid-owner "$CAPTURE_UID" -j REDIRECT --to-ports "$LISTEN_PORT"
iptables -t nat -C OUTPUT -p tcp -j "$CHAIN4" 2>/dev/null || iptables -t nat -A OUTPUT -p tcp -j "$CHAIN4"
if [[ "$ENABLE_UDP" == "1" ]]; then
  iptables -t nat -A "$CHAIN4" -p udp -m owner --uid-owner "$CAPTURE_UID" -j REDIRECT --to-ports "$LISTEN_PORT"
  iptables -t nat -C OUTPUT -p udp -j "$CHAIN4" 2>/dev/null || iptables -t nat -A OUTPUT -p udp -j "$CHAIN4"
fi

if command -v ip6tables >/dev/null 2>&1; then
  ip6tables -t nat -N "$CHAIN6" 2>/dev/null || true
  ip6tables -t nat -F "$CHAIN6"
  ip6tables -t nat -A "$CHAIN6" -d ::1/128 -j RETURN
  ip6tables -t nat -A "$CHAIN6" -m owner --uid-owner "$RUNTIME_USER" -j RETURN
  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
)
  ip6tables -t nat -A "$CHAIN6" -p tcp -m owner --uid-owner "$CAPTURE_UID" -j REDIRECT --to-ports "$LISTEN_PORT"
  ip6tables -t nat -C OUTPUT -p tcp -j "$CHAIN6" 2>/dev/null || ip6tables -t nat -A OUTPUT -p tcp -j "$CHAIN6"
  if [[ "$ENABLE_UDP" == "1" ]]; then
    ip6tables -t nat -A "$CHAIN6" -p udp -m owner --uid-owner "$CAPTURE_UID" -j REDIRECT --to-ports "$LISTEN_PORT"
    ip6tables -t nat -C OUTPUT -p udp -j "$CHAIN6" 2>/dev/null || ip6tables -t nat -A OUTPUT -p udp -j "$CHAIN6"
  fi
fi

echo "mynetspeeder transparent mode started on ${LISTEN_HOST}:${LISTEN_PORT}"
echo "capture uid: $CAPTURE_UID"
echo "udp capture: $ENABLE_UDP"
echo "log file: $LOG_FILE"

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