#!/bin/bash

# STUN 网络测试脚本 (纯ASCII + 中文)
# 用法: ./test_stun.sh [STUN_SERVER]
# 默认服务器: stunserver2025.stunprotocol.org

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# 默认服务器
STUN_SERVER="${1:-stunserver2025.stunprotocol.org}"
STUN_PORT="3478"

# 检查 stunclient 是否存在
if ! command -v stunclient &> /dev/null; then
    echo -e "${RED}[错误] 未找到 stunclient 命令${NC}"
    echo "请确保 stunclient 已编译并放置在 PATH 路径下"
    exit 1
fi

# 打印分隔线
print_separator() {
    echo -e "${CYAN}----------------------------------------${NC}"
}

# 打印测试标题
print_title() {
    echo -e "\n${BLUE}>>>> $1 ${NC}"
    print_separator
}

echo -e "${CYAN}+---------------------------------------+${NC}"
echo -e "${CYAN}|       STUN 网络环境测试脚本           |${NC}"
echo -e "${CYAN}+---------------------------------------+${NC}"
echo -e "目标服务器: ${YELLOW}$STUN_SERVER:$STUN_PORT${NC}"
echo -e "执行时间: $(date)"
echo ""

# ============================================================
# 测试1: 基础绑定测试 (UDP)
# ============================================================
print_title "测试1: 基础绑定测试 (UDP)"
echo -e "命令: stunclient $STUN_SERVER $STUN_PORT"
print_separator

RESULT=$(stunclient "$STUN_SERVER" "$STUN_PORT" 2>&1)

if echo "$RESULT" | grep -q "Binding test: success"; then
    LOCAL_ADDR=$(echo "$RESULT" | grep "Local address" | awk '{print $3}')
    MAPPED_ADDR=$(echo "$RESULT" | grep "Mapped address" | awk '{print $3}')
    echo -e "${GREEN}[通过] 绑定测试成功${NC}"
    echo -e "  内网地址: $LOCAL_ADDR"
    echo -e "  公网地址: $MAPPED_ADDR"
else
    echo -e "${RED}[失败] 绑定测试失败${NC}"
    echo "$RESULT"
fi

# ============================================================
# 测试2: NAT 行为测试 (Behavior)
# ============================================================
print_title "测试2: NAT 行为测试 (映射规则)"
echo -e "命令: stunclient --mode behavior $STUN_SERVER $STUN_PORT"
print_separator

RESULT=$(stunclient --mode behavior "$STUN_SERVER" "$STUN_PORT" 2>&1)

if echo "$RESULT" | grep -q "Behavior test: success"; then
    BEHAVIOR=$(echo "$RESULT" | grep "Nat behavior" | cut -d':' -f2- | xargs)
    echo -e "${GREEN}[通过] NAT 行为测试成功${NC}"
    echo -e "  映射类型: ${YELLOW}$BEHAVIOR${NC}"
    
    # 解读结果
    case "$BEHAVIOR" in
        *"Endpoint Independent Mapping"*)
            echo -e "  ${GREEN}-> 解读: 最宽松的锥形NAT，P2P打洞成功率最高${NC}"
            ;;
        *"Address Dependent Mapping"*)
            echo -e "  ${YELLOW}-> 解读: 地址依赖映射，P2P打洞成功率中等${NC}"
            ;;
        *"Address and Port Dependent Mapping"*)
            echo -e "  ${RED}-> 解读: 对称型NAT，P2P打洞困难，可能需要TURN中继${NC}"
            ;;
        *)
            echo -e "  ${CYAN}-> 解读: $BEHAVIOR${NC}"
            ;;
    esac
else
    echo -e "${RED}[失败] NAT 行为测试失败${NC}"
    echo "$RESULT"
fi

# ============================================================
# 测试3: NAT 过滤测试 (Filtering)
# ============================================================
print_title "测试3: NAT 过滤测试 (入站规则)"
echo -e "命令: stunclient --mode filtering --localport 9999 $STUN_SERVER $STUN_PORT"
print_separator

RESULT=$(stunclient --mode filtering --localport 9999 "$STUN_SERVER" "$STUN_PORT" 2>&1)

if echo "$RESULT" | grep -q "Filtering test: success"; then
    FILTERING=$(echo "$RESULT" | grep "Nat filtering" | cut -d':' -f2- | xargs)
    echo -e "${GREEN}[通过] NAT 过滤测试成功${NC}"
    echo -e "  过滤规则: ${YELLOW}$FILTERING${NC}"
    
    # 解读结果
    case "$FILTERING" in
        *"Endpoint Independent Filtering"*)
            echo -e "  ${GREEN}-> 解读: 最开放，任何外部主机均可连接（需先通信）${NC}"
            ;;
        *"Address Dependent Filtering"*)
            echo -e "  ${YELLOW}-> 解读: 中等严格，仅允许已通信的IP连接${NC}"
            ;;
        *"Address and Port Dependent Filtering"*)
            echo -e "  ${RED}-> 解读: 最严格，仅允许已通信的IP:端口连接${NC}"
            ;;
        *)
            echo -e "  ${CYAN}-> 解读: $FILTERING${NC}"
            ;;
    esac
else
    echo -e "${RED}[失败] NAT 过滤测试失败${NC}"
    echo "$RESULT"
fi

# ============================================================
# 测试4: TCP 协议测试
# ============================================================
print_title "测试4: TCP 协议测试"
echo -e "命令: stunclient --protocol tcp $STUN_SERVER $STUN_PORT"
print_separator

RESULT=$(stunclient --protocol tcp "$STUN_SERVER" "$STUN_PORT" 2>&1)

if echo "$RESULT" | grep -q "Binding test: success"; then
    LOCAL_ADDR=$(echo "$RESULT" | grep "Local address" | awk '{print $3}')
    MAPPED_ADDR=$(echo "$RESULT" | grep "Mapped address" | awk '{print $3}')
    echo -e "${GREEN}[通过] TCP 绑定测试成功${NC}"
    echo -e "  内网地址: $LOCAL_ADDR"
    echo -e "  公网地址: $MAPPED_ADDR"
else
    echo -e "${RED}[失败] TCP 绑定测试失败${NC}"
    echo "$RESULT"
fi

# ============================================================
# 测试总结
# ============================================================
print_title "测试总结"

echo -e "${CYAN}网络环境诊断结果:${NC}"
echo -e "  + 基础连接   : $(echo "$RESULT" | grep -q "success" && echo "${GREEN}正常${NC}" || echo "${RED}异常${NC}")"
echo -e "  + NAT 类型   : 请查看测试2结果"
echo -e "  + 过滤规则   : 请查看测试3结果"
echo -e "  + TCP 支持   : $(echo "$RESULT" | grep -q "success" && echo "${GREEN}支持${NC}" || echo "${YELLOW}不支持${NC}")"

echo ""
echo -e "${CYAN}建议:${NC}"
echo "  1. 如果NAT类型为 Endpoint Independent Mapping，P2P连接体验最佳"
echo "  2. 如果NAT类型为 Address and Port Dependent Mapping，建议配置TURN服务器保底"
echo "  3. 如果TCP测试失败，某些严格防火墙环境可能影响连接"

print_separator
echo -e "${GREEN}测试完成！${NC}"