#!/bin/bash

#--- To create queues with default params, adjust the values below and run script without params ---#

n=3				                #Number of users
IP_addr_3oct="12.0.0"	        #First three octets - Same for all IP addresses
IP_addr=(52 53 54 55 56 57)		#Just the last octect
total_rate=16384		        #Set the maximum Rate
indv_rate=$(( total_rate / n ))	#Initialize individual rates with equal values
device="eth3"			        #Interface

#--- HELP FUNCTION TO SHOW HOW TO ENTER PARAMTERS ---#
function usage 
{
	echo ""
	echo "(To create HTB with default values, run the script with no parameters)"
	echo "Create HTB:	                htb create -n [Number of users] -r [Total available rate kbps]"
	echo "Change rules:	                htb rules -ip [IP(last octect)] -min [Min Rate kbps] -max [Max Rate kbps]"
    echo "Kill queue:	                htb delete_child -ip [IP(last octect)]"
    echo "Change parent rule:           htb parent_rules -rate [New Rate kbps] -max [Max Rate kbps]"
	echo "Show Classes:	                htb show"
	echo "View per class bytes sent:    htb bytes_sent"
	echo "Reset HTB:	                htb reset"
	echo ""
}

#--- RESET FUNCTION FOR HBT. Remove all previouosly created queues. ---#
function reset
{
	echo "Deleting older queues..."
    #Deleting any previously created queues. (Will mention "no file exist" if no queue exists)
    tc qdisc del dev $device root
}

#--- CREATE FUNCTION FOR HBT. Bandwidth is divided equally among users at creation. ---#
function create
{

	echo "Deleting older queues..."
    #Deleting any previously created queues. (Will mention "no file exist" if no queue exists)
    tc qdisc del root dev $device
    
    # ensuring that TCP segmentation flag is set correctly in the interface
    # when we create the queues
    ethtool -K $device sg off
    ethtool -K $device tso off
    ethtool -K $device gso off
    ethtool -K $device gro off

    #Creating the root queue and class
    tc qdisc add dev $device root handle 1: htb
    echo ${total_rate}kbps
    tc class add dev $device parent 1: classid 1:1 htb rate $((${total_rate}*1000)) ceil $((${total_rate}*1000))

    # Creating a separate class for each User and adding a filter to assign users to queues.
    for (( i=0; i<n; i+=1 ))
    do
        ID=$(( ${IP_addr[i]}))
        # echo "tc class add dev $device parent 1:1 classid 1:$ID htb rate $((${indv_rate}*1000)) ceil $((${total_rate}*1000))"
        tc class add dev $device parent 1:1 classid 1:$ID htb rate $((${indv_rate}*1000)) ceil $((${total_rate}*1000))
        # echo "tc filter add dev $device protocol ip parent 1:0 prio 1 u32 match ip dst $IP_addr_3oct.${IP_addr[i]}/32 flowid 1:$ID"
        tc filter add dev $device protocol ip parent 1:0 prio 1 u32 match ip dst $IP_addr_3oct.${IP_addr[i]}/32 flowid 1:$ID
    done

    #Displaying the created classes based on IP addresses of users
    echo "Created Queues as follows..."
    tc class show dev $device
}


#--- FUNCTION TO CHANGE RULES. Have to specify ip, and minimum and maximum bandwdith. ---#
function rules
{
    echo "Updating child queue... 12.0.0."$ip
    ID=$(( $ip ))
    tc class change dev $device parent 1:1 classid 1:$ID htb rate $((${min_rate}*1000)) ceil $((${max_rate}*1000))
}

#--- FUNCTION TO DELETE QUEUE. Have to specify ip. ---#
function delete_child
{
    echo "Deleting child queue... 12.0.0."$ip "dst" ${IP_addr_3oct}.${ip}/32
    ID=$(( $ip ))
    # delete the filter first - getting error here, but seems fine
    tc filter del dev $device protocol ip parent 1:0 prio 1 u32 match ip dst ${IP_addr_3oct}.${ip}/32 flowid 1:${ID}
    # echo "tc filter del dev $device protocol ip parent 1:0 prio 1 u32 match ip dst ${IP_addr_3oct}.${ip}/32 flowid 1:${ID}"
    # then delete the class
    tc class del dev $device classid 1:$ID
    # echo "tc class del dev $device classid 1:$ID"
    
    # show the classes
    # tc class show dev $device
    echo 
}

#--- FUNCTION TO CHANGE RULES ON PARENT. Have to specify RATE AND CEIL bandwdiths. ---#
function parent_rule
{
    echo "Updating parent queue..."
    tc class change dev $device parent 1: classid 1:1 htb rate $((${new_rate}*1000)) ceil $((${max_rate}*1000))
}


#--- FUNCTION TO CREATE PARENT. Have to specify RATE bandwdith. ---#
function create_parent
{
    echo "Deleting older queues..."
    #Deleting any previously created queues. (Will mention "no file exist" if no queue exists)
    tc qdisc del root dev $device
    
    echo "Setting TCP flags..."
    # ensuring that TCP segmentation flag is set correctly in the interface
    # when we create the queues
    ethtool -K $device sg off
    ethtool -K $device tso off
    ethtool -K $device gso off
    ethtool -K $device gro off

    echo "Creating parent queue..."
    #Creating the root queue and class
    tc qdisc add dev $device root handle 1: htb
    echo ${total_rate}kbps
    tc class add dev $device parent 1: classid 1:1 htb rate $((${total_rate}*1000)) ceil $((${total_rate}*1000))
}



#--- ADJUSTING POSITIONAL PARAMTERES ---#
if [ "$1" == create ]; then
{
while [ "$2" != "" ]; do
    case $2 in
        -n | --users )           shift
                                n=$2
                                ;;
         -r | --rate )           shift
                                total_rate=$2
                                ;;
        -h | --help )           usage
                                exit
                                ;;
        * )                     usage
                                exit 1
    esac
    shift
done
indv_rate=$(( total_rate / n ))
create
}

# create a parent queue
elif [ "$1" == create_parent ]; then
{
while [ "$2" != "" ]; do
    case $2 in
         -r | --rate )           shift
                                total_rate=$2
                                ;;
        -h | --help )           usage
                                exit
                                ;;
        * )                     usage
                                exit 1
    esac
    shift
done
create_parent
}

elif [ "$1" == rules ]; then
{
while [ "$2" != "" ]; do
    case $2 in
        -ip | --IP )		shift
				ip=$2
				;;
         -min | --minrate )	shift
				min_rate=$2
				;;
         -max | --maxrate )	shift
				max_rate=$2
				;;
        -h | --help )           usage
                                exit
                                ;;
        * )                     usage
                                exit 1
    esac
    shift
done
rules
}

elif [ "$1" == delete_child ]; then
{
while [ "$2" != "" ]; do
    case $2 in
        -ip | --IP )		shift
				ip=$2
				;;
        -h | --help )           usage
                                exit
                                ;;
        * )                     usage
                                exit 1
    esac
    shift
done
delete_child
}

elif [ "$1" == parent_rule ]; then
{
while [ "$2" != "" ]; do
    case $2 in
         -rate | --newrate )	shift
				new_rate=$2
				;;
         -max | --maxrate )	shift
				max_rate=$2
				;;
         -h | --help )           usage
                                exit
                                ;;
        * )                     usage
                                exit 1
    esac
    shift
done
parent_rule
}

elif [ "$1" == show ]; then
{
    echo "TC show"
	tc class show dev $device
	echo "TC filter"
	tc filter show dev $device
}

elif [ "$1" == bytes_sent ]; then
{
	echo "TC list"
	# tc -s qdisc ls dev $device
	tc -s class show dev $device
}

elif [ "$1" == reset ]; then
{
	reset
}

elif [ "$1" == -h ] ||  [ "$1" == --help ]; then
{
	usage
}

else
#--- Creating Queues with default values ---#
	echo "Number of Users = $n"
	echo "Total Available Rate = ${total_rate}kbps"
	create

	exit
fi




