#!/usr/bin/env bash
# ============================================================
# check_stun.sh — STUN 服务器 UDP + TCP 双协议批量检测
# 依赖: bash, python3, nc (可选)
# 用法: bash check_stun.sh [-t 超时秒数] [-p 并发数]
# ============================================================

TIMEOUT=5
PARALLEL=5  # 并发协程数（加快大批量检测速度）

# ---- 颜色 ----
GREEN="\033[0;32m"; RED="\033[0;31m"; YELLOW="\033[1;33m"
CYAN="\033[0;36m"; RESET="\033[0m"; BOLD="\033[1m"

# ---- 参数解析 ----
while getopts "t:p:" opt; do
  case $opt in
    t) TIMEOUT="$OPTARG" ;;
    p) PARALLEL="$OPTARG" ;;
  esac
done

# ---- 服务器列表 ----
STUN_SERVERS=(
  "stun.12voip.com:3478"
  "stun.1und1.de:3478"
  "stun.callromania.ro:3478"
  "stun.chat.bilibili.com:3478"
  "stun.cloudflare.com:3478"
  "stun.easybell.de:3478"
  "stun.hc-yun.com:13478"
  "stun.l.google.com:19302"
  "stun.linphone.org:3478"
  "stun.miwifi.com:3478"
  "stun.nextcloud.com:443"
  "stun.nextcloud.com:3478"
  "stun.schlund.de:3478"
  "stun.antisip.com:3478"
  "stun.sonetel.net:3478"
  "stun.sonetel.com:3478"
  "stun.freeswitch.org:3478"
  "stun1.l.google.com:19302"
  "stun2.l.google.com:19302"
  "stun4.l.google.com:19302"
)

TOTAL=${#STUN_SERVERS[@]}
TMPDIR_WORK=$(mktemp -d)
trap 'rm -rf "$TMPDIR_WORK"' EXIT

# ============================================================
# Python3 并发检测核心（UDP STUN + TCP，asyncio）
# ============================================================
python3 - "$TIMEOUT" "$PARALLEL" "$TMPDIR_WORK" "${STUN_SERVERS[@]}" << 'PYEOF'
import sys, asyncio, struct, os, time

timeout   = int(sys.argv[1])
parallel  = int(sys.argv[2])
outdir    = sys.argv[3]
servers   = sys.argv[4:]

STUN_REQ = struct.pack('>HHI', 0x0001, 0, 0x2112A442) + os.urandom(12)

async def check_udp(host, port):
    t0 = time.monotonic()
    try:
        loop = asyncio.get_event_loop()
        transport, protocol = await asyncio.wait_for(
            loop.create_datagram_endpoint(
                lambda: UDPProto(),
                remote_addr=(host, port)
            ), timeout=timeout)
        transport.sendto(STUN_REQ)
        ok = await asyncio.wait_for(protocol.received.wait(), timeout=timeout)
        transport.close()
        ms = int((time.monotonic() - t0) * 1000)
        return ("ok", ms)
    except Exception:
        ms = int((time.monotonic() - t0) * 1000)
        return ("fail", ms)

class UDPProto(asyncio.DatagramProtocol):
    def __init__(self):
        self.received = asyncio.Event()
    def datagram_received(self, data, addr):
        self.received.set()
    def error_received(self, exc):
        self.received.set()

async def check_tcp(host, port):
    t0 = time.monotonic()
    try:
        _, writer = await asyncio.wait_for(
            asyncio.open_connection(host, port), timeout=timeout)
        writer.close()
        try: await writer.wait_closed()
        except: pass
        ms = int((time.monotonic() - t0) * 1000)
        return ("ok", ms)
    except Exception:
        ms = int((time.monotonic() - t0) * 1000)
        return ("fail", ms)

async def check_one(idx, entry):
    if ':' in entry:
        last_colon = entry.rfind(':')
        host = entry[:last_colon]
        port = int(entry[last_colon+1:])
    else:
        host = entry; port = 3478

    udp, tcp = await asyncio.gather(check_udp(host, port), check_tcp(host, port))
    result = f"{udp[0]}:{udp[1]}:{tcp[0]}:{tcp[1]}"
    with open(os.path.join(outdir, f"{idx:04d}"), "w") as f:
        f.write(result)
    # 实时打印进度到 stderr
    print(f"DONE:{idx}", file=sys.stderr, flush=True)

async def main():
    sem = asyncio.Semaphore(parallel)
    async def guarded(idx, entry):
        async with sem:
            await check_one(idx, entry)
    await asyncio.gather(*[guarded(i, s) for i, s in enumerate(servers)])

asyncio.run(main())
PYEOF

PY_EXIT=$?
if [[ $PY_EXIT -ne 0 ]]; then
  echo -e "${RED}错误: Python3 检测失败，请确认已安装 python3${RESET}"
  exit 1
fi

# ============================================================
# 汇总输出
# ============================================================
ok_udp=0; ok_tcp=0; fail_udp=0; fail_tcp=0
udp_ok_list=(); tcp_ok_list=()

fmt_ms() {
  local ms=$1
  if [[ $ms -ge 1000 ]]; then printf "%ds" $(( ms/1000 ))
  else printf "%dms" "$ms"; fi
}

col_status() {
  local s=$1 ms=$2
  local lat; lat=$(fmt_ms "$ms")
  if [[ "$s" == "ok" ]]; then printf "${GREEN}✔ OK${RESET}   (%-6s)" "$lat"
  else printf "${RED}✘ FAIL${RESET} (%-6s)" "$lat"; fi
}

echo ""
echo -e "${BOLD}╔════════════════════════════════════════════════════════════════════════╗${RESET}"
echo -e "${BOLD}║         STUN 服务器批量检测  UDP + TCP  (共 ${TOTAL} 条)              ║${RESET}"
echo -e "${BOLD}╚════════════════════════════════════════════════════════════════════════╝${RESET}"
echo -e "  超时: ${YELLOW}${TIMEOUT}s${RESET}   并发: ${YELLOW}${PARALLEL}${RESET}   UDP=STUN协议  TCP=端口探测"
echo ""
printf "  ${BOLD}%-40s  %-22s  %-22s${RESET}\n" "服务器" "UDP (STUN协议)" "TCP (端口探测)"
echo "  ──────────────────────────────────────────────────────────────────────────"

for i in "${!STUN_SERVERS[@]}"; do
  entry="${STUN_SERVERS[$i]}"
  result_file="$TMPDIR_WORK/$(printf '%04d' $i)"

  if [[ ! -f "$result_file" ]]; then
    printf "  %-40s  ${RED}(无结果)${RESET}\n" "$entry"
    continue
  fi

  IFS=: read -r us um ts tm < "$result_file"

  printf "  %-40s  " "$entry"
  col_status "$us" "$um"
  printf "  "
  col_status "$ts" "$tm"
  echo ""

  if [[ "$us" == "ok" ]]; then (( ok_udp++ ));  udp_ok_list+=("stun:${entry}"); else (( fail_udp++ )); fi
  if [[ "$ts" == "ok" ]]; then (( ok_tcp++ ));  tcp_ok_list+=("stun:${entry}"); else (( fail_tcp++ )); fi
done

echo "  ──────────────────────────────────────────────────────────────────────────"
echo ""
echo -e "  ${BOLD}汇总:${RESET}"
echo -e "    UDP  ${GREEN}✔ 可达: ${ok_udp}${RESET}   ${RED}✘ 不可达: ${fail_udp}${RESET}"
echo -e "    TCP  ${GREEN}✔ 可达: ${ok_tcp}${RESET}   ${RED}✘ 不可达: ${fail_tcp}${RESET}"
echo ""

if [[ ${#udp_ok_list[@]} -gt 0 ]]; then
  echo -e "  ${CYAN}${BOLD}UDP 可用列表 (推荐，STUN 标准协议):${RESET}"
  for r in "${udp_ok_list[@]}"; do echo -e "    ${GREEN}${r}${RESET}"; done
  echo ""
fi

if [[ ${#tcp_ok_list[@]} -gt 0 ]]; then
  echo -e "  ${CYAN}${BOLD}TCP 可用列表:${RESET}"
  for r in "${tcp_ok_list[@]}"; do echo -e "    ${GREEN}${r}${RESET}"; done
  echo ""
fi