#!/bin/bash

# Somewhat simple script that controls bandwidth usage.
#
# This script assumes that iptables, ebtables and tc are installed.
# It also assumes that it is being run on a bridge for the purposes
# of port matching.  This should be easy to modify.
# 
# By Matthew Strait, 2006.  May be distributed under the GPL version 2,
# http://www.gnu.org/licenses/gpl.txt

# tc needs to be told about the physical devices, even if you're a bridge
physdevs="eth0 eth1"

# syntax: "<match type> = <match arg> , <tc speed>".
# Match types are "layer7" and "port".
# "port" matches source or destination for tcp or udp.
# "kbps" means "KBytes/second".  This is tc's fault.
actions=( 
"port = 80, 20kbps"
"layer7 = gnutella, 50kbps"
"layer7 = dns, 500kbps"
"layer7 = smtp, 500kbps"
"layer7 = x11, 100kbps"
)

# Nothing below this line needs to be edited
#########################################################################

PATH=$PATH:/usr/local/sbin

# count commas...
lastaction=`echo ${actions[*]} | tr \, '\n' | wc -l`

# extra comma and array starts at zero...
let lastaction-=2

stop-tc() {
	if ! [ $1 ]; then
	        echo "specify a device!"
	        cleanup 1
	fi

	for dev in $@; do
	        if tc qdisc del dev $dev root &> /dev/null; then
	                echo "tc has now stopped for $dev"
	        else
	                echo "stopping tc for $dev failed (was probably already stopped)"
	        fi
	done
}

cleanup(){
	# Flush the whole filter table.
	iptables -F
	if ! [ $? = "0" ]; then 
		echo flushing iptables failed at line $LINENO
	else
		echo iptables flushed
	fi

	ebtables -F
	if ! [ $? = "0" ]; then 
		echo flushing ebtables failed at line $LINENO
	else
		echo ebtables flushed
	fi

	# stop traffic control completely.
	stop-tc $physdevs

	# if called with non-zero value, exit with that value
	if [[ $1 != 0 ]]; then exit $1; fi
}

# 0 = don't exit
cleanup 0

# If these aren't loaded manually, shaping of the child connections will not work.
# (Well, with version 0.1, it won't work anyway.)
for m in ip_conntrack_ftp ip_conntrack_irc ip_conntrack_tftp ip_conntrack_amanda; do
	if ! lsmod | grep $m > /dev/null; then
		if ! modprobe $m; then
			echo failed to load module $m
		fi
	fi
done

# set up basic traffic control magic
for dev in $physdevs; do
	tc qdisc add dev $dev root handle 1: htb default 10
	if ! [ $? = "0" ]; then echo tc failed at line $LINENO; cleanup 1; fi
done

# The Netfilter mark number and also the QoS queue number.
# Must start at 2 so it doesn't collide.
n=2

# index into the speeds array
index=0

for m in `seq 0 $lastaction`; do 

  match=`echo ${actions[$m]} | cut -d\, -f1`
  speed=`echo ${actions[$m]} | cut -d\, -f2`

  type=`echo $match | cut -d\= -f1`
  arg=`echo $match | cut -d\= -f2`

  echo Packets matching \"$match\" will be shaped to $speed \($n\)

  if [ $type == "layer7" ]; then
	echo $arg $n >> l7-tmp-conf.$$
  elif [ $type == "port" ]; then
	ebtables -A FORWARD -p IPv4 --ip-proto tcp --ip-source-port $arg -j mark --set-mark $n
	ebtables -A FORWARD -p IPv4 --ip-proto tcp --ip-destination-port $arg -j mark --set-mark $n
	ebtables -A FORWARD -p IPv4 --ip-proto udp --ip-source-port $arg -j mark --set-mark $n
	ebtables -A FORWARD -p IPv4 --ip-proto udp --ip-destination-port $arg -j mark --set-mark $n
  else
	echo "failed to parse \"$match\""
	cleanup 1
  fi

  if ! [ $? = "0" ]; then echo \[ip\|eb\]tables failed at line $LINENO; cleanup 1; fi

  for dev in $physdevs; do 
	  # per class traffic control black magic
	  tc class  add dev $dev parent 1:  classid 1:$n htb rate $speed burst 1k
	  if ! [ $? = "0" ]; then echo tc failed at line $LINENO; cleanup 1; fi

	  tc qdisc  add dev $dev parent 1:$n handle  $n: sfq perturb 10
	  if ! [ $? = "0" ]; then echo tc failed at line $LINENO; cleanup 1; fi

	  tc filter add dev $dev protocol ip parent 1: prio 1 handle $n fw flowid 1:$n
	  if ! [ $? = "0" ]; then echo tc failed at line $LINENO; cleanup 1; fi
  done

  let n++
  let index++
done

modprobe ip_conntrack_netlink
iptables -A FORWARD -j QUEUE
l7-filter -f l7-tmp-conf.$$

cleanup 0
rm l7-tmp-conf.$$