#!/bin/bash
#
# 网络流量监控脚本
# 功能: 实时监控网络接口流量，显示上传/下载速度和统计信息
# 兼容: RockyLinux, Ubuntu, Debian, CentOS 等主流 Linux 发行版
#

set -o pipefail

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

# 默认配置
INTERVAL=1
INTERFACE=""
COUNT=0  # 0 表示无限循环
SHOW_TOTAL=false
UNIT="auto"

# 使用说明
usage() {
    cat << EOF
用法: $0 [选项]

网络流量监控脚本 - 实时监控网络接口流量

选项:
  -i <interface>  指定网络接口 (默认: 自动检测活跃接口)
  -n <seconds>    刷新间隔秒数 (默认: 1)
  -c <count>      监控次数，0 为无限 (默认: 0)
  -t              显示累计流量统计
  -u <unit>       显示单位: auto, B, K, M, G (默认: auto)
  -l              列出所有可用网络接口
  -h              显示此帮助信息

示例:
  $0                      # 自动检测接口，实时监控
  $0 -i eth0              # 监控指定接口
  $0 -i eth0 -n 2         # 每 2 秒刷新一次
  $0 -i eth0 -c 10        # 监控 10 次后退出
  $0 -t                   # 显示累计流量统计
  $0 -l                   # 列出所有网络接口

快捷键:
  Ctrl+C  退出监控
EOF
    exit 0
}

# 错误退出
die() {
    echo -e "${RED}错误: $1${NC}" >&2
    exit 1
}

# 检查接口是否存在
check_interface() {
    local iface=$1
    if [ ! -d "/sys/class/net/${iface}" ]; then
        return 1
    fi
    return 0
}

# 获取所有网络接口
get_all_interfaces() {
    local interfaces=()
    for iface in /sys/class/net/*; do
        iface=$(basename "$iface")
        # 排除 lo 回环接口
        [ "$iface" = "lo" ] && continue
        interfaces+=("$iface")
    done
    echo "${interfaces[@]}"
}

# 自动检测活跃的网络接口
detect_active_interface() {
    local iface

    # 优先检测有默认路由的接口
    iface=$(ip route show default 2>/dev/null | awk '/default/ {print $5}' | head -1)
    if [ -n "$iface" ] && check_interface "$iface"; then
        echo "$iface"
        return 0
    fi

    # 否则选择第一个 UP 状态的非 lo 接口
    for iface in /sys/class/net/*; do
        iface=$(basename "$iface")
        [ "$iface" = "lo" ] && continue

        local state
        state=$(cat "/sys/class/net/${iface}/operstate" 2>/dev/null)
        if [ "$state" = "up" ]; then
            echo "$iface"
            return 0
        fi
    done

    # 最后选择第一个非 lo 接口
    for iface in /sys/class/net/*; do
        iface=$(basename "$iface")
        [ "$iface" = "lo" ] && continue
        echo "$iface"
        return 0
    done

    return 1
}

# 列出所有接口信息
list_interfaces() {
    echo ""
    echo -e "${CYAN}可用网络接口:${NC}"
    echo ""
    printf "  %-15s %-10s %-18s %s\n" "接口" "状态" "MAC 地址" "IP 地址"
    echo "  ---------------------------------------------------------------------------"

    for iface in /sys/class/net/*; do
        iface=$(basename "$iface")
        [ "$iface" = "lo" ] && continue

        local state mac ip
        state=$(cat "/sys/class/net/${iface}/operstate" 2>/dev/null || echo "unknown")
        mac=$(cat "/sys/class/net/${iface}/address" 2>/dev/null || echo "-")
        ip=$(ip addr show "$iface" 2>/dev/null | awk '/inet / {print $2}' | cut -d'/' -f1 | head -1)
        [ -z "$ip" ] && ip="-"

        local state_color
        if [ "$state" = "up" ]; then
            state_color="${GREEN}${state}${NC}"
        else
            state_color="${RED}${state}${NC}"
        fi

        printf "  %-15s %-10b %-18s %s\n" "$iface" "$state_color" "$mac" "$ip"
    done
    echo ""
}

# 获取接口流量数据
# 返回: rx_bytes tx_bytes
get_traffic() {
    local iface=$1
    local rx_bytes tx_bytes

    rx_bytes=$(cat "/sys/class/net/${iface}/statistics/rx_bytes" 2>/dev/null)
    tx_bytes=$(cat "/sys/class/net/${iface}/statistics/tx_bytes" 2>/dev/null)

    if [ -z "$rx_bytes" ] || [ -z "$tx_bytes" ]; then
        return 1
    fi

    echo "$rx_bytes $tx_bytes"
}

# 格式化字节数
format_bytes() {
    local bytes=$1
    local unit=$2
    local result

    case "$unit" in
        B)
            result="${bytes} B"
            ;;
        K)
            result=$(awk "BEGIN {printf \"%.2f KB\", $bytes / 1024}")
            ;;
        M)
            result=$(awk "BEGIN {printf \"%.2f MB\", $bytes / 1024 / 1024}")
            ;;
        G)
            result=$(awk "BEGIN {printf \"%.2f GB\", $bytes / 1024 / 1024 / 1024}")
            ;;
        auto|*)
            if [ "$bytes" -lt 1024 ]; then
                result="${bytes} B"
            elif [ "$bytes" -lt 1048576 ]; then
                result=$(awk "BEGIN {printf \"%.2f KB\", $bytes / 1024}")
            elif [ "$bytes" -lt 1073741824 ]; then
                result=$(awk "BEGIN {printf \"%.2f MB\", $bytes / 1024 / 1024}")
            else
                result=$(awk "BEGIN {printf \"%.2f GB\", $bytes / 1024 / 1024 / 1024}")
            fi
            ;;
    esac

    echo "$result"
}

# 格式化速率
format_rate() {
    local bytes_per_sec=$1
    local unit=$2
    local result

    case "$unit" in
        B)
            result="${bytes_per_sec} B/s"
            ;;
        K)
            result=$(awk "BEGIN {printf \"%.2f KB/s\", $bytes_per_sec / 1024}")
            ;;
        M)
            result=$(awk "BEGIN {printf \"%.2f MB/s\", $bytes_per_sec / 1024 / 1024}")
            ;;
        G)
            result=$(awk "BEGIN {printf \"%.2f GB/s\", $bytes_per_sec / 1024 / 1024 / 1024}")
            ;;
        auto|*)
            if [ "$bytes_per_sec" -lt 1024 ]; then
                result="${bytes_per_sec} B/s"
            elif [ "$bytes_per_sec" -lt 1048576 ]; then
                result=$(awk "BEGIN {printf \"%.2f KB/s\", $bytes_per_sec / 1024}")
            elif [ "$bytes_per_sec" -lt 1073741824 ]; then
                result=$(awk "BEGIN {printf \"%.2f MB/s\", $bytes_per_sec / 1024 / 1024}")
            else
                result=$(awk "BEGIN {printf \"%.2f GB/s\", $bytes_per_sec / 1024 / 1024 / 1024}")
            fi
            ;;
    esac

    echo "$result"
}

# 显示累计流量统计
show_total_stats() {
    local interfaces

    if [ -n "$INTERFACE" ]; then
        interfaces=("$INTERFACE")
    else
        read -ra interfaces <<< "$(get_all_interfaces)"
    fi

    echo ""
    echo -e "${CYAN}网络流量统计${NC}"
    echo "时间: $(date '+%Y-%m-%d %H:%M:%S')"
    echo ""
    printf "  %-15s %15s %15s %15s %15s\n" "接口" "接收流量" "发送流量" "接收包数" "发送包数"
    echo "  ---------------------------------------------------------------------------------"

    local total_rx=0 total_tx=0

    for iface in "${interfaces[@]}"; do
        if ! check_interface "$iface"; then
            continue
        fi

        local rx_bytes tx_bytes rx_packets tx_packets
        rx_bytes=$(cat "/sys/class/net/${iface}/statistics/rx_bytes" 2>/dev/null || echo 0)
        tx_bytes=$(cat "/sys/class/net/${iface}/statistics/tx_bytes" 2>/dev/null || echo 0)
        rx_packets=$(cat "/sys/class/net/${iface}/statistics/rx_packets" 2>/dev/null || echo 0)
        tx_packets=$(cat "/sys/class/net/${iface}/statistics/tx_packets" 2>/dev/null || echo 0)

        total_rx=$((total_rx + rx_bytes))
        total_tx=$((total_tx + tx_bytes))

        printf "  %-15s %15s %15s %15s %15s\n" \
            "$iface" \
            "$(format_bytes "$rx_bytes" "$UNIT")" \
            "$(format_bytes "$tx_bytes" "$UNIT")" \
            "$rx_packets" \
            "$tx_packets"
    done

    echo "  ---------------------------------------------------------------------------------"
    printf "  %-15s %15s %15s\n" "总计" "$(format_bytes "$total_rx" "$UNIT")" "$(format_bytes "$total_tx" "$UNIT")"
    echo ""
}

# 清除屏幕并移动光标
clear_screen() {
    # 使用 ANSI 转义序列清屏并移动光标到顶部
    printf "\033[2J\033[H"
}

# 打印监控头部
print_monitor_header() {
    local iface=$1
    local state ip mac

    state=$(cat "/sys/class/net/${iface}/operstate" 2>/dev/null || echo "unknown")
    ip=$(ip addr show "$iface" 2>/dev/null | awk '/inet / {print $2}' | cut -d'/' -f1 | head -1)
    mac=$(cat "/sys/class/net/${iface}/address" 2>/dev/null || echo "-")

    echo "================================================================================"
    echo -e "                      ${CYAN}网络流量实时监控${NC}"
    echo "================================================================================"
    echo ""
    echo -e "  接口: ${BOLD}${iface}${NC}"
    echo "  状态: ${state}"
    echo "  MAC:  ${mac}"
    echo "  IP:   ${ip:-无}"
    echo "  时间: $(date '+%Y-%m-%d %H:%M:%S')"
    echo ""
    echo "--------------------------------------------------------------------------------"
    printf "  %-12s %18s %18s %18s\n" "" "下载 ↓" "上传 ↑" "总计"
    echo "--------------------------------------------------------------------------------"
}

# 实时监控
monitor_realtime() {
    local iface=$1
    local prev_rx=0 prev_tx=0
    local curr_rx curr_tx
    local rx_rate tx_rate
    local session_rx=0 session_tx=0
    local iteration=0

    # 获取初始值
    read -r prev_rx prev_tx <<< "$(get_traffic "$iface")"

    # 捕获 Ctrl+C
    trap 'echo ""; echo -e "${YELLOW}监控已停止${NC}"; exit 0' INT

    while true; do
        sleep "$INTERVAL"

        # 获取当前值
        read -r curr_rx curr_tx <<< "$(get_traffic "$iface")"

        if [ -z "$curr_rx" ] || [ -z "$curr_tx" ]; then
            echo -e "${RED}获取流量数据失败${NC}"
            continue
        fi

        # 计算速率
        rx_rate=$(( (curr_rx - prev_rx) / INTERVAL ))
        tx_rate=$(( (curr_tx - prev_tx) / INTERVAL ))

        # 处理计数器溢出（32位系统）
        if [ "$rx_rate" -lt 0 ]; then
            rx_rate=0
        fi
        if [ "$tx_rate" -lt 0 ]; then
            tx_rate=0
        fi

        # 累计本次会话流量
        session_rx=$((session_rx + (curr_rx - prev_rx)))
        session_tx=$((session_tx + (curr_tx - prev_tx)))

        # 清屏并打印
        clear_screen
        print_monitor_header "$iface"

        # 当前速率
        printf "  ${GREEN}当前速率${NC}   %18s %18s\n" \
            "$(format_rate "$rx_rate" "$UNIT")" \
            "$(format_rate "$tx_rate" "$UNIT")"

        # 本次会话流量
        printf "  ${BLUE}本次会话${NC}   %18s %18s %18s\n" \
            "$(format_bytes "$session_rx" "$UNIT")" \
            "$(format_bytes "$session_tx" "$UNIT")" \
            "$(format_bytes "$((session_rx + session_tx))" "$UNIT")"

        # 累计流量
        printf "  ${CYAN}累计流量${NC}   %18s %18s %18s\n" \
            "$(format_bytes "$curr_rx" "$UNIT")" \
            "$(format_bytes "$curr_tx" "$UNIT")" \
            "$(format_bytes "$((curr_rx + curr_tx))" "$UNIT")"

        echo "--------------------------------------------------------------------------------"
        echo ""
        echo -e "  刷新间隔: ${INTERVAL} 秒 | 按 ${YELLOW}Ctrl+C${NC} 退出"

        # 更新上一次的值
        prev_rx=$curr_rx
        prev_tx=$curr_tx

        # 检查是否达到指定次数
        ((iteration++))
        if [ "$COUNT" -gt 0 ] && [ "$iteration" -ge "$COUNT" ]; then
            echo ""
            echo -e "${GREEN}已完成 $COUNT 次监控${NC}"
            break
        fi
    done
}

# 解析参数
while getopts "i:n:c:u:tlh" opt; do
    case $opt in
        i)
            INTERFACE="$OPTARG"
            ;;
        n)
            INTERVAL="$OPTARG"
            if ! [[ "$INTERVAL" =~ ^[0-9]+$ ]] || [ "$INTERVAL" -lt 1 ]; then
                die "刷新间隔必须是大于 0 的整数: $INTERVAL"
            fi
            ;;
        c)
            COUNT="$OPTARG"
            if ! [[ "$COUNT" =~ ^[0-9]+$ ]]; then
                die "监控次数必须是非负整数: $COUNT"
            fi
            ;;
        u)
            UNIT="$OPTARG"
            case "$UNIT" in
                auto|B|K|M|G) ;;
                *)
                    die "无效的单位: $UNIT (可选: auto, B, K, M, G)"
                    ;;
            esac
            ;;
        t)
            SHOW_TOTAL=true
            ;;
        l)
            list_interfaces
            exit 0
            ;;
        h)
            usage
            ;;
        *)
            usage
            ;;
    esac
done

# 如果只显示累计流量
if [ "$SHOW_TOTAL" = true ]; then
    show_total_stats
    exit 0
fi

# 自动检测接口
if [ -z "$INTERFACE" ]; then
    INTERFACE=$(detect_active_interface)
    if [ -z "$INTERFACE" ]; then
        die "未找到可用的网络接口"
    fi
fi

# 检查接口是否存在
if ! check_interface "$INTERFACE"; then
    die "网络接口 '$INTERFACE' 不存在"
fi

# 开始监控
monitor_realtime "$INTERFACE"