#!/bin/bash

# 定义网络接口和限速参数
INTERFACE="eno1"
LIMIT_RATE="1mbit"
PORTS=("8337" "8338")  # 指定限速的端口

# 加载 ifb 模块并确保 ifb0 存在且启用
prepare_ifb() {
    modprobe ifb 2>/dev/null
    ip link show ifb0 >/dev/null 2>&1 || ip link add ifb0 type ifb
    ip link set ifb0 up 2>/dev/null
}

# 清理旧规则
clear_rules() {
    echo "Removing old bandwidth rules..."
    tc qdisc del dev "$INTERFACE" root 2>/dev/null || true
    tc qdisc del dev "$INTERFACE" ingress 2>/dev/null || true
    tc qdisc del dev ifb0 root 2>/dev/null || true
    ip link set ifb0 down 2>/dev/null || true
    echo "Old rules removed successfully."
}

# 添加所有端口的限速
add_all_ports_limit() {
    echo "Applying bandwidth limit to all ports..."

    # 添加上行规则
    tc qdisc add dev "$INTERFACE" root handle 1: htb default 11
    tc class add dev "$INTERFACE" parent 1: classid 1:1 htb rate "$LIMIT_RATE" ceil "$LIMIT_RATE"
    tc class add dev "$INTERFACE" parent 1:1 classid 1:11 htb rate "$LIMIT_RATE" ceil "$LIMIT_RATE"

    # 未指定过滤器，默认限制所有流量
    echo "All ports are now limited to $LIMIT_RATE."

    # 添加下行规则
    tc qdisc add dev ifb0 root handle 1: htb default 11
    tc class add dev ifb0 parent 1: classid 1:1 htb rate "$LIMIT_RATE" ceil "$LIMIT_RATE"
    tc class add dev ifb0 parent 1:1 classid 1:11 htb rate "$LIMIT_RATE" ceil "$LIMIT_RATE"

    # 将 ingress 流量重定向到 ifb0
    tc qdisc add dev "$INTERFACE" ingress
    tc filter add dev "$INTERFACE" parent ffff: protocol ip prio 1 u32 match ip protocol 0x06 0xff action mirred egress redirect dev ifb0

    echo "Bandwidth limit applied to all ports successfully."
}

# 添加指定端口的限速
add_specified_ports_limit() {
    echo "Applying bandwidth limit to specified ports (${PORTS[*]})..."

    # 添加上行规则
    tc qdisc add dev "$INTERFACE" root handle 1: htb
    tc class add dev "$INTERFACE" parent 1: classid 1:1 htb rate "$LIMIT_RATE" ceil "$LIMIT_RATE"

    for PORT in "${PORTS[@]}"; do
        tc filter add dev "$INTERFACE" protocol ip parent 1:0 prio 1 u32 match ip dport "$PORT" 0xffff flowid 1:1
    done

    # 添加下行规则
    tc qdisc add dev ifb0 root handle 1: htb
    tc class add dev ifb0 parent 1: classid 1:1 htb rate "$LIMIT_RATE" ceil "$LIMIT_RATE"

    for PORT in "${PORTS[@]}"; do
        tc filter add dev ifb0 protocol ip parent 1:0 prio 1 u32 match ip sport "$PORT" 0xffff flowid 1:1
    done

    # 将 ingress 流量重定向到 ifb0
    tc qdisc add dev "$INTERFACE" ingress
    for PORT in "${PORTS[@]}"; do
        tc filter add dev "$INTERFACE" parent ffff: protocol ip prio 1 u32 match ip sport "$PORT" 0xffff action mirred egress redirect dev ifb0
    done

    echo "Bandwidth limit applied to specified ports (${PORTS[*]}) successfully."
}

case "$1" in
    addall)
        clear_rules
        prepare_ifb
        add_all_ports_limit
        ;;
    add)
        clear_rules
        prepare_ifb
        add_specified_ports_limit
        ;;
    del)
        clear_rules
        ;;
    *)
        echo "Usage: $0 {addall|add|del}"
        exit 1
        ;;
esac