#!/bin/sh
SNAME=`basename $0`

FW="/sbin/iptables"
TC="/sbin/tc"

P2P_MARK="10"
LOW_MARK="11"
IFACES="eth0 eth1"

LINK_RATE="3000kbit"
P2P_RATE="128kbit"
LOW_RATE="128kbit"

ROOT="1"
NORMAL="2"
P2P="3"
LOW="4"

if [ ! -e $FW ]; then
    echo "$SNAME: Error: Couldn't find iptables binary."
    exit 1
fi

if [ ! -e $TC ]; then
    echo "$SNAME: Error: Couldn't find tc binary."
    exit 1
fi

function start () {
    echo -n "Setting p2p networks limit... "
    $FW -t mangle -L PREROUTING | grep ipp2p > /dev/null
    if [ $? -ne 0 ]; then
        $FW -t mangle -A PREROUTING -j CONNMARK --restore-mark
        $FW -t mangle -A PREROUTING -m mark ! --mark 0 -j ACCEPT
        $FW -t mangle -A PREROUTING -m ipp2p --ipp2p --bit --apple --winmx --soul -j MARK --set-mark $P2P_MARK
        $FW -t mangle -A PREROUTING -m mark ! --mark 0 -j CONNMARK --save-mark

#       $FW -t mangle -A PREROUTING -j CONNMARK --restore-mark
#       $FW -t mangle -A PREROUTING -m mark ! --mark 0 -j ACCEPT
#       $FW -t mangle -A PREROUTING -d 192.168.90.153 -j MARK --set-mark $LOW_MARK
#       $FW -t mangle -A PREROUTING -s 192.168.90.153 -j MARK --set-mark $LOW_MARK
#       $FW -t mangle -A PREROUTING -m mark ! --mark 0 -j CONNMARK --save-mark
    fi
    for i in $IFACES; do
        $TC qdisc add dev $i root handle 1: htb default $NORMAL
        $TC class add dev $i parent 1: classid 1:$ROOT htb rate $LINK_RATE ceil $LINK_RATE
        $TC class add dev $i parent 1:$ROOT classid 1:$NORMAL htb rate $LINK_RATE ceil $LINK_RATE
        $TC class add dev $i parent 1:$ROOT classid 1:$P2P htb rate $P2P_RATE ceil $P2P_RATE
        $TC class add dev $i parent 1:$ROOT classid 1:$LOW htb rate $LOW_RATE ceil $LOW_RATE
        $TC filter add dev $i parent 1:0 protocol ip prio 2 handle $P2P_MARK fw classid 1:$P2P
        $TC filter add dev $i parent 1:0 protocol ip prio 1 handle $LOW_MARK fw classid 1:$LOW
    done
    echo "done."
}

function stop () {
    echo -n "Removing p2p networks limit... "
    for i in $IFACES; do
        $TC qdisc del dev $i root
    done
    echo "done."
}

function restart () {
    stop
    start
}

function status () {
    for i in $IFACES; do
        $TC -s class show dev $i
    done
}

case "$1" in
    start)
        start
        ;;
    stop)
        stop
        ;;
    restart)
        restart
        ;;
    status)
        status
        ;;
    *)
        echo "Usage: $SNAME (start|stop|restart|status)"
        exit 1
esac

exit 0