#!/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 (IPv4) TO TCP (IPv4)"
  echo "2) TCP (IPv6) TO TCP (IPv6)"
  echo "3) UDP (IPv4) TO UDP (IPv4)"
  echo "4) UDP (IPv6) TO UDP (IPv6)"
  echo "5) TCP (IPv4) TO TCP (IPv6)"
  echo "6) TCP (IPv6) TO TCP (IPv4)"
  echo "7) UDP (IPv4) TO UDP (IPv6)"
  echo "8) UDP (IPv6) TO UDP (IPv4)"

  read -rp "选择 (1/2/3/4/5/6/7/8): " PROTO_CHOICE

  case "$PROTO_CHOICE" in
    1)
      PROTO="tcp"
      SOCAT_CMD="TCP-LISTEN:${LOCAL_PORT},reuseaddr,fork TCP:${TARGET_IP}:${TARGET_PORT}"
      ;;
    2)
      PROTO="tcp6"
      SOCAT_CMD="TCP6-LISTEN:${LOCAL_PORT},reuseaddr,fork TCP6:${TARGET_IP}:${TARGET_PORT}"
      ;;
    3)
      PROTO="udp"
      SOCAT_CMD="UDP-LISTEN:${LOCAL_PORT},reuseaddr,fork UDP:${TARGET_IP}:${TARGET_PORT}"
      ;;
    4)
      PROTO="udp6"
      SOCAT_CMD="UDP6-LISTEN:${LOCAL_PORT},reuseaddr,fork UDP6:${TARGET_IP}:${TARGET_PORT}"
      ;;
    5)
      PROTO="tcp4to6"
      SOCAT_CMD="TCP6-LISTEN:${LOCAL_PORT},reuseaddr,fork TCP:${TARGET_IP}:${TARGET_PORT}"
      ;;
    6)
      PROTO="tcp6to4" 
      SOCAT_CMD="TCP-LISTEN:${LOCAL_PORT},reuseaddr,fork TCP6:${TARGET_IP}:${TARGET_PORT}"
      ;;
    7)
      PROTO="udp4to6"
      SOCAT_CMD="UDP6-LISTEN:${LOCAL_PORT},reuseaddr,fork UDP:${TARGET_IP}:${TARGET_PORT}"
      ;;
    8)
      PROTO="udp6to4"
      SOCAT_CMD="UDP-LISTEN:${LOCAL_PORT},reuseaddr,fork UDP6:${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