#!/bin/bash
#
# 端口连通性检测脚本
# 功能: 批量检测多台服务器的指定端口连通性
#

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

# 默认配置
TIMEOUT=3
HOSTS_FILE=""
SUCCESS_COUNT=0
FAIL_COUNT=0
FAILED_LIST=()

# 使用说明
usage() {
    echo "用法: $0 [选项] [主机 端口]"
    echo ""
    echo "选项:"
    echo "  -f <file>     从文件读取服务器列表"
    echo "  -t <seconds>  连接超时时间 (默认: 3秒)"
    echo "  -h            显示此帮助信息"
    echo ""
    echo "示例:"
    echo "  $0 192.168.1.100 22"
    echo "  $0 192.168.1.100 22,80,443"
    echo "  $0 -f hosts.txt"
    echo "  $0 -t 5 192.168.1.100 22"
    echo ""
    echo "配置文件格式 (每行一条记录):"
    echo "  192.168.1.100 22,80"
    echo "  web.example.com 80,443"
    exit 1
}

# 分隔线
print_line() {
    echo "================================================================================"
}

print_dash_line() {
    echo "--------------------------------------------------------------------------------"
}

# 头部信息
print_header() {
    print_line
    echo -e "                         ${CYAN}端口连通性检测${NC}"
    echo "                    $(date '+%Y-%m-%d %H:%M:%S')"
    print_line
    echo ""
}

# 检测单个端口
check_port() {
    local host=$1
    local port=$2
    local start_time end_time latency status

    start_time=$(date +%s%3N)

    # 优先使用 nc，否则使用 /dev/tcp
    if command -v nc &> /dev/null; then
        nc -z -w "$TIMEOUT" "$host" "$port" &> /dev/null
        status=$?
    else
        timeout "$TIMEOUT" bash -c "echo >/dev/tcp/$host/$port" &> /dev/null
        status=$?
    fi

    end_time=$(date +%s%3N)
    latency=$((end_time - start_time))

    if [ $status -eq 0 ]; then
        printf " %-25s %-9s ${GREEN}✓ 成功${NC}     %dms\n" "$host" "$port" "$latency"
        ((SUCCESS_COUNT++))
    else
        printf " %-25s %-9s ${RED}✗ 失败${NC}     -\n" "$host" "$port"
        ((FAIL_COUNT++))
        FAILED_LIST+=("$host:$port")
    fi
}

# 检测主机的多个端口
check_host() {
    local host=$1
    local ports=$2

    IFS=',' read -ra PORT_ARRAY <<< "$ports"
    for port in "${PORT_ARRAY[@]}"; do
        port=$(echo "$port" | tr -d ' ')
        [ -n "$port" ] && check_port "$host" "$port"
    done
}

# 从文件读取并检测
check_from_file() {
    local file=$1
    local total_hosts=0
    local total_ports=0

    if [ ! -f "$file" ]; then
        echo -e "${RED}错误: 文件 '$file' 不存在${NC}"
        exit 1
    fi

    # 统计总数
    while IFS= read -r line || [ -n "$line" ]; do
        line=$(echo "$line" | sed 's/#.*//' | xargs)
        [ -z "$line" ] && continue

        local host=$(echo "$line" | awk '{print $1}')
        local ports=$(echo "$line" | awk '{print $2}')

        [ -z "$host" ] || [ -z "$ports" ] && continue

        ((total_hosts++))
        IFS=',' read -ra PORT_ARRAY <<< "$ports"
        total_ports=$((total_ports + ${#PORT_ARRAY[@]}))
    done < "$file"

    echo "检测目标: $total_hosts 台服务器, $total_ports 个端口"
    echo "超时时间: ${TIMEOUT} 秒"
    echo ""
    print_dash_line
    printf " %-25s %-9s %-10s %s\n" "主机" "端口" "状态" "延迟"
    print_dash_line

    # 执行检测
    while IFS= read -r line || [ -n "$line" ]; do
        line=$(echo "$line" | sed 's/#.*//' | xargs)
        [ -z "$line" ] && continue

        local host=$(echo "$line" | awk '{print $1}')
        local ports=$(echo "$line" | awk '{print $2}')

        [ -z "$host" ] || [ -z "$ports" ] && continue

        check_host "$host" "$ports"
    done < "$file"

    print_dash_line
}

# 打印结果汇总
print_summary() {
    local total=$((SUCCESS_COUNT + FAIL_COUNT))
    local success_rate=0
    local fail_rate=0

    if [ $total -gt 0 ]; then
        success_rate=$((SUCCESS_COUNT * 100 / total))
        fail_rate=$((100 - success_rate))
    fi

    echo ""
    print_line
    echo -e "                           ${CYAN}检测结果汇总${NC}"
    print_line
    echo " 总计检测: $total"
    echo -e " 成功: ${GREEN}$SUCCESS_COUNT${NC} (${success_rate}%)"
    echo -e " 失败: ${RED}$FAIL_COUNT${NC} (${fail_rate}%)"

    if [ ${#FAILED_LIST[@]} -gt 0 ]; then
        echo ""
        echo -e " ${RED}失败列表:${NC}"
        for item in "${FAILED_LIST[@]}"; do
            echo "   - $item"
        done
    fi
    print_line
}

# 解析参数
while getopts "f:t:h" opt; do
    case $opt in
        f)
            HOSTS_FILE="$OPTARG"
            ;;
        t)
            TIMEOUT="$OPTARG"
            ;;
        h)
            usage
            ;;
        *)
            usage
            ;;
    esac
done

shift $((OPTIND - 1))

# 主逻辑
print_header

if [ -n "$HOSTS_FILE" ]; then
    check_from_file "$HOSTS_FILE"
elif [ $# -ge 2 ]; then
    echo "检测目标: $1"
    echo "检测端口: $2"
    echo "超时时间: ${TIMEOUT} 秒"
    echo ""
    print_dash_line
    printf " %-25s %-9s %-10s %s\n" "主机" "端口" "状态" "延迟"
    print_dash_line
    check_host "$1" "$2"
    print_dash_line
else
    usage
fi

print_summary