#!/bin/bash
# @desc: TCP & UDP 延迟测试客户端
# @usage: wcl udping_client <目标IP> [端口]
#   例: wcl udping_client 1.1.1.1 443
# ==============================================================================
# Script Name: TCP & UDP Latency Tester (Final Complete)
# Author:      LULU
# Description: 双协议(TCP/UDP)延迟测试客户端。
#              集成 DNS 解析、握手延迟计算、丢包率统计及完整的安装指南。
# ==============================================================================

HOST=$1
PORT=${2:-9999}

# --- [核心功能] 帮助与安装指南 (已补全) ---
show_help() {
    echo "=================================================================="
    echo "   TCP & UDP Latency Tester - 使用手册"
    echo "=================================================================="
    echo "用法: $0 <目标IP或域名> [端口]"
    echo "示例: $0 p2p1.melulu.top 10022"
    echo ""
    echo "测试项目:"
    echo "  1. TCP 连接握手耗时 (Connect Latency)"
    echo "  2. TCP 数据往返延迟 (Ping-Pong RTT)"
    echo "  3. UDP 数据往返延迟 (Ping-Pong RTT)"
    echo "  4. UDP 丢包率统计"
    echo ""
    echo "支持的工具 (脚本会自动检测并优先使用):"
    echo "  1. Nping   (Nmap组件 - 推荐, 毫秒级精度)"
    echo "  2. Python3 (推荐, 通用性最强)"
    echo "  3. Netcat  (仅做基本连通性检查)"
    echo ""
    echo "------------------------------------------------------------------"
    echo "   各平台工具安装指南 (客户端建议安装 Nmap)"
    echo "------------------------------------------------------------------"
    echo "[OpenWrt / 嵌入式Linux]"
    echo "  opkg update"
    echo "  opkg install nmap-nping     # (首选, 仅安装nping)"
    echo "  opkg install python3-light  # (备选)"
    echo ""
    echo "[Ubuntu / Debian / Kali]"
    echo "  sudo apt update"
    echo "  sudo apt install nmap python3"
    echo ""
    echo "[CentOS / RHEL / Fedora]"
    echo "  sudo yum install nmap python3"
    echo ""
    echo "[Alpine Linux (Docker)]"
    echo "  apk add nmap python3"
    echo "=================================================================="
}

# 检查参数
if [ -z "$HOST" ] || [[ "$1" == "-h" ]] || [[ "$1" == "--help" ]]; then
    show_help
    exit 1
fi

echo "=============================================="
echo "   Testing -> $HOST:$PORT"
echo "=============================================="

# --- 方案 1: Nping (最专业) ---
if command -v nping >/dev/null 2>&1; then
    echo "[Tool] 使用 Nping (High Precision)"
    
    echo ""
    echo "--- [1/2] TCP 测试 ---"
    # --tcp-connect: 进行完整的三次握手来测试延迟
    nping --tcp-connect -p $PORT -c 5 "$HOST" | grep -E "Avg rtt|RCVD|Max rtt"
    
    echo ""
    echo "--- [2/2] UDP 测试 ---"
    # --source-port: 固定源端口 55555 以辅助穿透部分 NAT/防火墙
    nping --udp -p $PORT --source-port 55555 -c 5 --data-length 16 "$HOST" \
        | grep -v "ICMP" | grep -v "unreachable" | grep -E "Avg rtt|RCVD|Max rtt|Lost"
    
    exit 0
fi

# --- 方案 2: Python3 (最通用) ---
if command -v python3 >/dev/null 2>&1; then
    echo "[Tool] 使用 Python3 (Native)"
    
    python3 -u -c "
import socket, time, sys

target = '$HOST'
port = int($PORT)

# --- DNS 解析模块 ---
try:
    print(f'[Info] Resolving {target}...')
    ip = socket.gethostbyname(target)
    print(f'[Info] Target IP: {ip}')
except Exception as e:
    print(f'[Error] DNS解析失败: {e}')
    print('提示: 请检查域名拼写或 DNS 设置')
    sys.exit(1)

# --- TCP 测试模块 ---
print('\n--- [1/2] TCP 测试 ---')
tcp_rtt = []
try:
    # 1. 测试连接握手耗时
    t1 = time.time()
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(3) # 3秒连接超时
    s.connect((ip, port))
    t2 = time.time()
    print(f'TCP Connect (Handshake): {(t2-t1)*1000:.2f} ms')
    
    # 2. 测试数据传输 RTT
    for i in range(5):
        try:
            start = time.time()
            s.send(b'ping')
            data = s.recv(1024)
            end = time.time()
            if data:
                rtt = (end - start) * 1000
                tcp_rtt.append(rtt)
                print(f'Seq={i+1} RTT={rtt:.2f} ms')
        except Exception as e:
            print(f'Seq={i+1} Error: {e}')
        time.sleep(0.5)
    s.close()
    
    if tcp_rtt:
        print(f'TCP Avg RTT: {sum(tcp_rtt)/len(tcp_rtt):.2f} ms')
except Exception as e:
    print(f'[TCP Error] 连接失败: {e}')
    print('提示: 请检查服务端防火墙是否放行 TCP 端口')

# --- UDP 测试模块 ---
print('\n--- [2/2] UDP 测试 ---')
udp_rtt = []
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.settimeout(2) # 2秒接收超时

for i in range(5):
    try:
        start = time.time()
        s.sendto(b'ping', (ip, port))
        data, _ = s.recvfrom(1024)
        end = time.time()
        rtt = (end - start) * 1000
        udp_rtt.append(rtt)
        print(f'Seq={i+1} RTT={rtt:.2f} ms')
    except socket.timeout:
        print(f'Seq={i+1} Timeout')
    except Exception as e:
        print(f'Seq={i+1} Error: {e}')
    time.sleep(0.5)

if udp_rtt:
    print(f'UDP Avg RTT: {sum(udp_rtt)/len(udp_rtt):.2f} ms')
    print(f'丢包率: {100 - (len(udp_rtt)/5.0)*100:.1f}%')
else:
    print('UDP 丢包率: 100%')
    print('提示: 请检查服务端防火墙 UDP 端口，或云服务器安全组设置')
"
    exit 0
fi

# --- 方案 3: Netcat (兜底降级) ---
if command -v nc >/dev/null 2>&1; then
    echo "[Tool] 使用 Netcat (仅连通性测试)"
    echo "[Warn] 当前环境缺少 Nping 或 Python，无法进行高精度 RTT 测量。"
    
    echo ""
    echo "--- [1/2] TCP 测试 ---"
    # -z 扫描模式, -v 详细, -w 超时
    if nc -z -v -w 2 "$HOST" "$PORT"; then
        echo "TCP Connection: OK"
    else
        echo "TCP Connection: FAILED"
    fi
    
    echo ""
    echo "--- [2/2] UDP 测试 ---"
    echo "尝试发送 UDP 数据包..."
    # 使用 time 命令粗略估算
    for i in {1..3}; do
        (time echo "ping" | nc -u -w 1 "$HOST" "$PORT") 2>&1 | grep "real"
    done
    exit 0
fi

echo "[Error] 未找到可用工具 (Nping/Python/Nc)。"
echo "请运行 $0 --help 查看安装指南。"
exit 1