#!/bin/bash

set -e

SERVICE_PREFIX="socat-"
SYSTEMD_DIR="/etc/systemd/system"

############################
# Root 权限检查
############################
if [ "$EUID" -ne 0 ]; then
  echo "❌ 请使用 root 用户运行此脚本"
  exit 1
fi

############################
# 检查并可选安装 socat
############################
install_socat() {
  if command -v socat >/dev/null 2>&1; then
    return
  fi

  echo "⚠️ 未检测到 socat"
  read -rp "是否安装 socat 并继续？(y/N): " CONFIRM
  [[ "$CONFIRM" =~ ^[Yy]$ ]] || exit 1

  . /etc/os-release

  case "$ID" in
    ubuntu|debian)
      apt update -y && apt install -y socat
      ;;
    centos|rhel|almalinux|rocky)
      yum install -y socat
      ;;
    fedora)
      dnf install -y socat
      ;;
    arch)
      pacman -Sy --noconfirm socat
      ;;
    opensuse*|sles)
      zypper install -y socat
      ;;
    *)
      echo "❌ 不支持的系统，请手动安装 socat"
      exit 1
      ;;
  esac
}

install_socat

############################
# 列出 socat systemd 服务
############################
list_services() {
  echo
  echo "📋 当前 socat 转发规则："
  echo "--------------------------------------------------"

  mapfile -t SERVICES < <(systemctl list-unit-files \
    | awk '{print $1}' \
    | grep "^${SERVICE_PREFIX}.*\.service" || true)

  if [ "${#SERVICES[@]}" -eq 0 ]; then
    echo "（暂无 socat 规则）"
    return 1
  fi

  for i in "${!SERVICES[@]}"; do
    STATUS=$(systemctl is-active "${SERVICES[$i]}" 2>/dev/null || echo "unknown")
    printf "%2d) %-30s [%s]\n" "$((i+1))" "${SERVICES[$i]}" "$STATUS"
  done

  return 0
}

############################
# 创建新规则
############################
create_service() {
  echo
  read -rp "本地监听端口: " LOCAL_PORT
  read -rp "目标 IP 地址: " TARGET_IP
  read -rp "目标端口: " TARGET_PORT

  echo "协议类型："
  echo "1) TCP"
  echo "2) UDP"
  read -rp "选择 (1/2): " PROTO_CHOICE

  case "$PROTO_CHOICE" in
    1)
      PROTO="tcp"
      SOCAT_CMD="TCP-LISTEN:${LOCAL_PORT},reuseaddr,fork TCP:${TARGET_IP}:${TARGET_PORT}"
      ;;
    2)
      PROTO="udp"
      SOCAT_CMD="UDP-LISTEN:${LOCAL_PORT},reuseaddr UDP:${TARGET_IP}:${TARGET_PORT}"
      ;;
    *)
      echo "❌ 无效选择"
      return
      ;;
  esac

  SERVICE_NAME="${SERVICE_PREFIX}${PROTO}-${LOCAL_PORT}.service"
  SERVICE_FILE="${SYSTEMD_DIR}/${SERVICE_NAME}"

  cat > "$SERVICE_FILE" <<EOF
[Unit]
Description=Socat ${PROTO^^} Port Forward ${LOCAL_PORT} -> ${TARGET_IP}:${TARGET_PORT}
After=network.target

[Service]
Type=simple
ExecStart=/usr/bin/socat ${SOCAT_CMD}
Restart=always
RestartSec=3
LimitNOFILE=1048576

[Install]
WantedBy=multi-user.target
EOF

  systemctl daemon-reload
  systemctl enable "$SERVICE_NAME" --now

  echo "✅ 已创建并启动：$SERVICE_NAME"
}

############################
# 删除规则
############################
delete_service() {
  list_services || return

  echo
  read -rp "请输入要删除的规则编号: " INDEX
  [[ "$INDEX" =~ ^[0-9]+$ ]] || { echo "❌ 输入无效"; return; }

  SERVICE="${SERVICES[$((INDEX-1))]}"
  [[ -n "$SERVICE" ]] || { echo "❌ 编号不存在"; return; }

  echo "⚠️ 即将删除：$SERVICE"
  read -rp "确认删除？(y/N): " CONFIRM
  [[ "$CONFIRM" =~ ^[Yy]$ ]] || return

  systemctl stop "$SERVICE" 2>/dev/null || true
  systemctl disable "$SERVICE" 2>/dev/null || true
  rm -f "${SYSTEMD_DIR}/${SERVICE}"

  systemctl daemon-reload
  systemctl reset-failed

  echo "🗑️ 已删除：$SERVICE"
}

############################
# 主菜单
############################
while true; do
  echo
  echo "========== Socat 转发规则管理 =========="
  echo "1) 创建新的转发规则"
  echo "2) 查看已有转发规则"
  echo "3) 删除转发规则"
  echo "0) 退出"
  echo "======================================="
  read -rp "请选择: " CHOICE

  case "$CHOICE" in
    1) create_service ;;
    2) list_services ;;
    3) delete_service ;;
    0) exit 0 ;;
    *) echo "❌ 无效选择" ;;
  esac
done