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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(dirname "$(dirname "$SCRIPT_DIR")")"
PACKAGE_PARENT="$(dirname "$ROOT_DIR")"
PACKAGE_NAME="$(basename "$ROOT_DIR")"

source "$SCRIPT_DIR/runtime-lib.sh"

CONFIG_PATH="${1:-$ROOT_DIR/config-tcp.json}"
RUNTIME_USER="${MYNETSPEEDER_TCP_USER:-mynetspeeder}"
PID_FILE="/var/run/mynetspeeder-tcp-edge.pid"
LOG_FILE="/var/log/mynetspeeder-tcp-edge.log"
CHAIN4="MYNETSPEEDER_TCP"
CHAIN6="MYNETSPEEDER_TCP6"
SSH_PORTS="${MYNETSPEEDER_TCP_SSH_PORTS:-}"
CAPTURE_UID="${MYNETSPEEDER_TCP_CAPTURE_UID:-}"
SELF_EXCLUDE_V4="127.0.0.0/8 169.254.0.0/16"
SELF_EXCLUDE_V6="::1/128 fe80::/10"

require_root

if [[ ! -f "$CONFIG_PATH" ]]; then
  echo "config not found: $CONFIG_PATH"
  exit 1
fi

LISTEN_HOST_FROM_CONFIG="$(python3 - <<'PY' "$CONFIG_PATH"
import json, sys
cfg = json.load(open(sys.argv[1]))
print(cfg.get("listen_host", "127.0.0.1"))
PY
)"
LISTEN_PORT_FROM_CONFIG="$(python3 - <<'PY' "$CONFIG_PATH"
import json, sys
cfg = json.load(open(sys.argv[1]))
port = int(cfg.get("listen_port", 19080) or 19080)
print(port if port > 0 else 19080)
PY
)"
LISTEN_HOST="${MYNETSPEEDER_TCP_LISTEN_HOST:-$LISTEN_HOST_FROM_CONFIG}"
LISTEN_PORT="${MYNETSPEEDER_TCP_LISTEN_PORT:-$LISTEN_PORT_FROM_CONFIG}"

if ! [[ "$LISTEN_PORT" =~ ^[0-9]+$ ]]; then
  echo "listen port must be numeric"
  exit 1
fi

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

ensure_runtime_user "$RUNTIME_USER"
ensure_log_file "$LOG_FILE" "$RUNTIME_USER"
stop_pid_file "$PID_FILE" "edge-tcp --listen-host ${LISTEN_HOST} --listen-port ${LISTEN_PORT} --config ${CONFIG_PATH}"

cleanup_rules() {
  ensure_nat_jump_absent iptables nat OUTPUT -p tcp -j "$CHAIN4"
  ensure_nat_chain_absent iptables "$CHAIN4"
  ensure_nat_jump_absent ip6tables nat OUTPUT -p tcp -j "$CHAIN6"
  ensure_nat_chain_absent ip6tables "$CHAIN6"
}

cleanup_required=1
cleanup_done=0

cleanup_on_exit() {
  if [[ "$cleanup_done" -eq 1 || "$cleanup_required" -eq 0 ]]; then
    return 0
  fi
  cleanup_done=1
  cleanup_rules
  stop_pid_file "$PID_FILE" "edge-tcp --listen-host ${LISTEN_HOST} --listen-port ${LISTEN_PORT} --config ${CONFIG_PATH}"
}

trap 'cleanup_on_exit' ERR EXIT INT TERM

readarray -t RELAY_HOSTS < <(python3 - <<'PY' "$CONFIG_PATH"
import json, sys
cfg = json.load(open(sys.argv[1]))
for relay in cfg.get("relays", []):
    print(relay["host"])
PY
)

ARGV_JSON="$(python3 - <<'PY' "$PACKAGE_NAME" "$LISTEN_HOST" "$LISTEN_PORT" "$CONFIG_PATH"
import json, sys
package_name, listen_host, listen_port, config_path = sys.argv[1:]
print(json.dumps([
    "python3", "-m", package_name, "edge-tcp",
    "--listen-host", listen_host,
    "--listen-port", listen_port,
    "--config", config_path,
]))
PY
)"

start_python_service "$RUNTIME_USER" "$PACKAGE_PARENT" "$PACKAGE_PARENT" "$LOG_FILE" "$PID_FILE" \
  "edge-tcp --listen-host ${LISTEN_HOST} --listen-port ${LISTEN_PORT} --config ${CONFIG_PATH}" "$ARGV_JSON" >/dev/null

if ! wait_for_tcp_listen "$LISTEN_HOST" "$LISTEN_PORT" 15; then
  tail -n 50 "$LOG_FILE" || true
  exit 1
fi

iptables -t nat -N "$CHAIN4" 2>/dev/null || true
iptables -t nat -F "$CHAIN4"
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
for host in "${RELAY_HOSTS[@]}"; do
  [[ -n "$host" && "$host" != *:* ]] && iptables -t nat -A "$CHAIN4" -d "$host" -j RETURN
done
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_nat_jump_absent iptables nat OUTPUT -p tcp -j "$CHAIN4"
iptables -t nat -A OUTPUT -p tcp -j "$CHAIN4"

if command -v ip6tables >/dev/null 2>&1 && ip6tables -t nat -S >/dev/null 2>&1; then
  ip6tables -t nat -N "$CHAIN6" 2>/dev/null || true
  ip6tables -t nat -F "$CHAIN6"
  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
  for host in "${RELAY_HOSTS[@]}"; do
    [[ -n "$host" && "$host" == *:* ]] && ip6tables -t nat -A "$CHAIN6" -d "$host" -j RETURN
  done
  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_nat_jump_absent ip6tables nat OUTPUT -p tcp -j "$CHAIN6"
  ip6tables -t nat -A OUTPUT -p tcp -j "$CHAIN6"
fi

cleanup_required=0
trap - ERR EXIT INT TERM
echo "tcp-only started on ${LISTEN_HOST}:${LISTEN_PORT}"