#!/bin/bash

# wshaper.hfsc -- H-FSC based traffic shaper script, based on the idea
#                 and implementation found on flo.xssn.at by Florian Pritz
# Copyleft (C) 2009 Adrian C. <anrxc_sysphere_org>
#               All Rights Reversed

# Name is a play on wondershaper (lartc.org) as I already had wshaper.cbq
# and wshaper.htb scripts on my system. For more information on H-FSC see:
#   - http://www.cs.cmu.edu/~hzhang/HFSC/main.html


########################################################################
# Network control
# 
#   B/w throttle in ideal conditions:
#     4096*0.88, 4096*0.85
#       3604.48,   3481.60
# 
#      256*0.88,  256*0.85
#        225.28,    217.60
#
#   ADSL, actual b/w around: 3440kbps / 256kbps
# 
# D/l 4096kbps
DOWNLINK=3384
#
# U/l 256kbps
UPLINK=217
#
# I/face
DEV=ppp0


########################################################################
# Traffic segmentation
#
# IP's of VoIP phones, if any
VOIPIPS=""
#
# VoIP telephony, Skype, Ventrilo etc.
VOIPPORTS="3784 7977"
#
# Interactive: SSH, DNS, Dbox and gaming (ET, OA)
INTERACTIVEPORTS="22 53 12039 27950 27960 27965 28785 28786 28952"
#
# Web traffic, Jabber and IRC
BROWSINGPORTS="80 443 6667 6697 5222 5223 8080 9050"
#
# Everything unspecified will fall between Web and Data
#N/A
#
# Data transfers: FTP, Mail and Rsync
DATAPORTS="21 25 110 143 465 873 993 995"
#
# Lowest priority traffic: Bittorrent and other P2P traffic
P2PPORTS="6881:6999 10311:10325"


########################################################################
# Start
#
function check_device() {
    if [ -z "$DEV" ] ; then
        echo "$0: no interface specified"
        exit -1
    fi
}


function stop() {
    check_device

    # Reset everything to a known state (cleared)
    tc qdisc del dev $DEV root    &> /dev/null
    tc qdisc del dev $DEV ingress &> /dev/null
    #
    # Flush and delete tables
    iptables -t mangle --delete       POSTROUTING -o $DEV -j THESHAPER &> /dev/null
    iptables -t mangle --flush        THESHAPER &> /dev/null
    iptables -t mangle --delete-chain THESHAPER &> /dev/null

    echo "Shaping removed on interface: $DEV"
}


function start() {
    check_device

    if [ -z "$DOWNLINK" ] ; then
	echo "$0: no interface specified"
	exit -1
    fi
    if [ -z "$UPLINK" ] ; then
	echo "$0: no interface specified"
	exit -1
    fi


    # Traffic classes:
    #   1:2 Interactive (SSH, DNS, ACK, Games)
    #   1:3 Low latency (VoIP, Skype)
    #   1:4 Browsing (HTTP, IM, IRC)
    #   1:5 Default
    #   1:6 Middle-low priority (data)
    #   1:7 Lowest priority (p2p)

    # Install root HFSC qdisc
    tc qdisc add dev $DEV root handle 1: hfsc default 5

    # Add main rate limit class
    tc class add dev $DEV parent 1: classid 1:1 hfsc \
        sc rate ${UPLINK}kbit ul rate ${UPLINK}kbit

    # Interactive traffic: guarantee realtime full uplink for 50ms, then 5/10 of the uplink
    tc class add dev $DEV parent 1:1  classid 1:2 hfsc \
        rt m1   ${UPLINK}kbit d  50ms m2 $((5*$UPLINK/10))kbit \
        ls m1   ${UPLINK}kbit d  50ms m2 $((7*$UPLINK/10))kbit \
        ul rate ${UPLINK}kbit

    # VoIP: guarantee full uplink for 200ms, then 3/10
    tc class add dev $DEV parent 1:1  classid 1:3 hfsc \
        sc m1 ${UPLINK}kbit d 200ms m2 $((3*$UPLINK/10))kbit \
        ul rate ${UPLINK}kbit

    # Browsing: guarantee 3/10 uplink for 200ms, then guarantee 1/10
    tc class add dev $DEV parent 1:1  classid 1:4 hfsc \
        sc m1 $((3*$UPLINK/10))kbit d 200ms m2 $((1*$UPLINK/10))kbit \
        ul rate ${UPLINK}kbit

    # Default traffic: don't guarantee anything for the first second, then guarantee 1/10
    tc class add dev $DEV parent 1:1  classid 1:5 hfsc \
        sc m1         0 d    1s m2 $((1*$UPLINK/10))kbit \
        ul rate ${UPLINK}kbit

    # Middle-low taffic: don't guarantee anything for the first 5 seconds, then guarantee 1/10
    tc class add dev $DEV parent 1:1  classid 1:6 hfsc \
        sc m1         0 d   5s m2 $((1*$UPLINK/10))kbit \
        ul rate ${UPLINK}kbit

    # Lowest taffic: don't guarantee anything for the first 10 seconds,then guarantee 1/20
    tc class add dev $DEV parent 1:1  classid 1:7 hfsc \
        sc m1         0 d   10s m2 $((1*$UPLINK/20))kbit \
        ul rate ${UPLINK}kbit

    # Add THESHAPER chain to the mangle table in iptables
    iptables -t mangle --new-chain THESHAPER
    iptables -t mangle --insert    POSTROUTING -o $DEV -j THESHAPER

    # To speed up downloads while an upload is going on, put short ACK packets in the interactive class
    iptables -t mangle -A THESHAPER \
        -p tcp \
        -m tcp --tcp-flags FIN,SYN,RST,ACK ACK \
        -m length --length :64 \
        -j CLASSIFY --set-class 1:2

    # Put large (512+) icmp packets in browsing category
    iptables -t mangle -A THESHAPER \
        -p icmp \
        -m length --length 512: \
        -j CLASSIFY --set-class 1:4

    # ICMP (ip protocol 1) in the interactive class
    iptables -t mangle -A THESHAPER \
        -p icmp  \
        -m length --length :512 \
        -j CLASSIFY --set-class 1:2


    # Classify our traffic ports
    setclassbyport() {
        port=$1
        CLASS=$2
        iptables -t mangle -A THESHAPER -p udp --sport $port -j CLASSIFY --set-class $CLASS
        iptables -t mangle -A THESHAPER -p udp --dport $port -j CLASSIFY --set-class $CLASS
        iptables -t mangle -A THESHAPER -p tcp --sport $port -j CLASSIFY --set-class $CLASS
        iptables -t mangle -A THESHAPER -p tcp --dport $port -j CLASSIFY --set-class $CLASS
    }
    for port in $INTERACTIVEPORTS;  do setclassbyport $port 1:2; done
    for port in $VOIPPORTS;         do setclassbyport $port 1:3; done
    for port in $BROWSINGPORTS;     do setclassbyport $port 1:4; done
    for port in $DATAPORTS;         do setclassbyport $port 1:6; done
    for port in $P2PPORTS;          do setclassbyport $port 1:7; done
    #
    # Classify VoIP phones, if any
    for VOIP in $VOIPIPS
    do
        iptables -t mangle -A THESHAPER --src $VOIP -j CLASSIFY --set-class 1:3
        iptables -t mangle -A THESHAPER --dst $VOIP -j CLASSIFY --set-class 1:3
    done


    # Try to control the incoming traffic as well
    #
    # Set up ingress qdisc
    tc qdisc add dev $DEV handle ffff: ingress


    # Filter everything that is coming in too fast
    #
    # It's mostly HTTP downloads that keep jamming the downlink, try to restrict them to 95/100 of d/l
    tc filter add dev $DEV parent ffff: protocol ip prio 50 \
        u32 match ip src 0.0.0.0/0 \
        match ip protocol 6 0xff \
        match ip sport 80 0xffff \
        police rate $((95*${DOWNLINK}/100))kbit \
        burst 10k drop flowid :1

    tc filter add dev $DEV parent ffff: protocol ip prio 50 \
        u32 match ip src 0.0.0.0/0 \
        match ip protocol 6 0xff \
        match ip dport 80 0xffff \
        police rate $((95*${DOWNLINK}/100))kbit \
        burst 10k drop flowid :1

    echo "Shaping started on interface: $DEV"
}


function status() {
    check_device

    echo -e "\n\n\tTraffic statistics for interface: $DEV\n"

    #echo -e "[bandwidth]"
    #    vnstat -s | grep today

    echo -e "\n[qdisc]"
        tc -s qdisc show dev $DEV

    echo -e "\n[class]"
        tc -s class show dev $DEV

    echo -e "\n[filter]"
        tc -s filter show dev $DEV

    echo -e "\n[iptables]"
        iptables -t mangle -L THESHAPER -v -x 2> /dev/null
}


case "$1" in
    status)
        status
    ;;
    stop) 
        stop 
    ;;
    start) 
        start
    ;;
    restart)
        stop
        start
    ;;
    *)
        echo "$0 {start|stop|status|restart} [iface]"
        exit
    ;;
esac