eos-select-bandwidth.py

#!/usr/bin/python2

# Script inspired from https://gist.github.com/obscurerichard/3740206

import os
import sys
import subprocess
import argparse
import textwrap

# Profiles taken from facebook's ATC project
# https://github.com/facebook/augmented-traffic-control/tree/master/utils/profiles
# 'name' : [ bandwidth, latency, packetloss ],
# 'default' key taken from average data from a few people in Rio de Janeiro 
network_type = {
    # values in megabit/s , miliseconds, and packetloss %
    'default'    :  [],
    '2G-rural'   :  [ 0.020, 0.018,  650,  2 ],
    '2G-urban'   :  [ 0.035, 0.032,  650,  2 ],
    '3G-avg'     :  [ 0.780, 0.330,  100,  1 ],
    '3G-good'    :  [ 0.850, 0.420,   90,  1 ],
    'cable'      :  [ 6.000, 1.000,    2,  1 ],
    'DSL'        :  [ 2.000, 0.256,    5,  1 ],
    'EDGE-avg'   :  [ 0.240, 0.200,  440,  1 ],
    'EDGE-good'  :  [ 0.250, 0.200,  360,  1 ],
    'EDGE-lossy' :  [ 0.240, 0.200,  420,  2 ]
}

def setParser():
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
                                        description=textwrap.dedent('''\
    Simulates a low bandwidth, high-latency network connection.
    Requires a Linux operating system with the \'tc\' traffic control tool.
    Available NETWORK profiles are:\n\t{0}'''.format('\n\t'.join( sorted(network_type))) ),
                                        epilog="Feel the pain our users feel.")
    parser.add_argument("interface",
                    help="Set network interface on which the changes will be applied; i.e. eth0, wlan0, etc")
    parser.add_argument("-o", default='slow', dest="command", type=str, choices=['slow', 'reset', 'status'],
                    help="Set operation to run. Defaults to 'slow'. \n")
    parser.add_argument("-n", default='default', dest="network",
                    help="Sets NETWORK profile from the available list (see --help)." )
    parser.add_argument("-down", default=1.2, dest="ingress", type=float,
                    help="Set max download bandwidth in megabits/s, number specified will be applied with a -20%% range")
    parser.add_argument("-up", default=0.6, dest="egress", type=float,
                    help="Set max upload bandwidth in megabits/s, number specified will be applied with a -20%% range")
    parser.add_argument("-l", default=330, dest="latency", type=int,
                    help="Set latency, this added on top of your real latency")
    parser.add_argument("-p", default=3, dest="packetloss", type=float,
                    help="Set packet loss this is a percentage value")
    return parser.parse_args()

def runCommand(cmd, out):
    if out == 0:
        return_code = subprocess.call(cmd, shell=True)
        return return_code
    elif out == 1:
        output = subprocess.check_output(cmd, shell=True)
        return output
    raise

def slow(iface, network):
    reset(iface)
    ingress_rate = network[0]
    ingress_ceil = ingress_rate + ( ingress_rate * 0.2)
    egress_rate = network[1]
    egress_ceil = egress_rate + ( egress_rate * 0.2)
    latency = network[2]
    packetloss = network[3]
    setup = [
        # INGRESS STUFF
        # load module
        "modprobe ifb",
        # enable ifb0 interface
        "ip link set dev ifb0 up",
        # raise ingress QDISC, redirect traffic from physical interface to ifb
        "tc qdisc add dev {0} handle ffff: ingress".format(iface),
        "tc filter add dev {0} parent ffff: protocol ip u32 match u32 0 0 action mirred egress redirect dev ifb0".format(iface),
        "tc qdisc add dev ifb0 root handle 1: htb default 10",
        "tc class add dev ifb0 parent 1:1 classid 1:10 htb rate {0}mbit ceil {1}mbit".format(ingress_rate, ingress_ceil),
        "tc qdisc add dev ifb0 parent 1:10 netem delay {0}ms loss {1}%".format(latency, packetloss),
        # EGRESS STUFF
        # raise egress QDISC and rules
        "tc qdisc add dev {0} root handle 1: htb default 12".format(iface),
        "tc class add dev {0} parent 1:1 classid 1:12 htb rate {1}mbit ceil {2}mbit".format(iface, egress_rate, egress_ceil),
    ]
    for step in setup:
        print step
        if runCommand(step, 0) != 0:
            print "Unable to raise queues interface configurations. Cleaning up."
            reset()
            sys.exit(1)
    print "All rules successfully applied."

def reset(iface):
    print "\nClearing rules...\n"
    teardown = [
        "tc qdisc del dev {0} root".format(iface),
        "tc qdisc del dev {0} ingress".format(iface),
        "tc qdisc del dev ifb0 root",
        "ip link set dev ifb0 down",
        "modprobe -r ifb"
    ]
    for step in teardown:
        runCommand(step, 0)
    print "Done."

def status():
    print "\nCurrent rules:\n"
    cmd = "tc qdisc"
    output = subprocess.check_output(cmd, shell=True)
    print output

def main(argv):
    args = setParser()
    if not os.getuid() == 0:
        print "You need to run this script as sudo"
        sys.exit(1)
    if len(sys.argv)==1:
        print "\nNetwork interface needs to be specified.\n"
        parser.print_help()
        sys.exit(1)
    if not args.network in network_type.keys():
        print "Available network profiles are:\n{0}".format('\t'+'\n\t'.join( sorted(network_type)))
        sys.exit(1)
    if args.command == 'status':
        status()
    elif args.command == 'reset':
        reset(args.interface)
    else:
        network_type['default'] = [args.ingress, args.egress, args.latency, args.packetloss]
        slow(args.interface, network_type[args.network])
    sys.exit(1)

if __name__ == "__main__":
    main(sys.argv)