#!/usr/bin/python3
# -*- coding: UTF-8 -*-
"""
DO NOT MODIFY THIS FILE. YOUR CHANGES WILL BE OVERWRITTEN
To extend this functionality add a new hook.
Add, update, and remove tc qdiscs and classes for Quality of Service.
"""
# © 2012, Multapplied Networks, Inc.

import sys
import os
import json
import subprocess

QOS_DIRECTORY = os.environ.get('BONDING_QOS_DIRECTORY', '/var/lib/bonding/qos')
QOS_CONFIG = os.path.join(QOS_DIRECTORY, 'tun%s.json')
TC_CMD = '/sbin/tc'


def popen(args, ignore_errors=False):
    p = subprocess.Popen(args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    p.wait()
    if not ignore_errors and p.returncode != 0:
        sys.exit("Failed to run QoS command: %s\n    ::%s" % (' '.join(args), p.stderr.read()))


def get_config_data(tunnel_id):
    qos_config_file = QOS_CONFIG % tunnel_id
    try:
        with open(qos_config_file) as f:
            config = json.loads(f.read())
    except IOError as e:
        sys.exit('Error opening file %s: %s' % (qos_config_file, e))
    except ValueError:
        sys.exit('Could not decode QoS file %s as JSON' % qos_config_file)
    traffic_classes = config['traffic_classes']
    traffic_classes.sort(key=lambda k: k['order'])
    return (traffic_classes, config['overhead_margin'])


def apply_rules(traffic_classes, overhead_margin, mode):
    speed = int(os.environ['SPEED']) * (1 - (overhead_margin / 100.0))

    major_id = traffic_classes[0]['profile']

    # Add the default tc rule first
    for tc in traffic_classes:
        if tc['default']:
            minor_id = 0xff + tc['id']
            if mode == 'add':
                popen([TC_CMD, 'qdisc', mode, 'dev', os.environ['INTERFACE'], 'root', 'handle', '%x:' % major_id, 'htb', 'default', '%x' % minor_id])
            popen([TC_CMD, 'class', mode, 'dev', os.environ['INTERFACE'], 'parent', '%x:' % major_id, 'classid', '%x:1' % major_id, 'htb', 'rate', '%skbit' % speed, 'ceil', '%skbit' % speed])
            break

    # Add the rest of the tc rules
    remaining_bandwidth = speed
    remaining_percent = 100.0
    for tc in traffic_classes:
        # Determine reserved amount
        if remaining_bandwidth <= 0:
            class_reserved = 10  # Give them just a little bit of reserved
            # bandwidth so the system doesn't fall apart.
            # Unless they have a minimum set.
            if tc['reserved_min'] is not None:
                class_reserved = tc['reserved_min'] * 1000
        else:
            initial_remaining = remaining_percent
            class_percent = tc['reserved_percent'] / remaining_percent
            class_reserved = remaining_bandwidth * class_percent

            if tc['reserved_max'] is not None:
                class_reserved = min(class_reserved, tc['reserved_max'] * 1000)
            if tc['reserved_min'] is not None:
                class_reserved = max(class_reserved, tc['reserved_min'] * 1000)

            remaining_bandwidth -= class_reserved
            remaining_percent -= class_percent * initial_remaining

        # Determine limit amount
        class_limit = None
        if any([tc['limit_percent'] is not None, tc['limit_max'] is not None, tc['limit_min'] is not None]):
            if tc['limit_percent'] is not None:
                class_limit = speed * (tc['limit_percent'] / 100.0)
            if tc['limit_max'] is not None:
                class_limit = min(class_limit, tc['limit_max'] * 1000)
            if tc['limit_min'] is not None:
                class_limit = max(class_limit, tc['limit_min'] * 1000)

        minor_id = 0xff + tc['id']
        class_args = [TC_CMD, 'class', mode, 'dev', os.environ['INTERFACE'], 'parent', '%x:1' % major_id, 'classid', '%x:%x' % (major_id, minor_id), 'htb', 'rate', '%skbit' % class_reserved]

        if class_limit is not None:
            class_args += ['ceil', '%skbit' % class_limit]
        else:
            class_args += ['ceil', '%skbit' % speed]
        class_args += ['prio', str(tc['order'])]
        popen(class_args)

        if mode == 'add':
            if tc['leaf_qdisc'] == 'sfq':
                qdisc = ['sfq', 'perturb', '10']
            else:
                qdisc = [tc['leaf_qdisc']]
            popen([TC_CMD, 'qdisc', mode, 'dev', os.environ['INTERFACE'], 'parent', '%x:%x' % (major_id, minor_id)] + qdisc)


def start(config):
    stop(ignore_errors=True)
    apply_rules(config[0], config[1], 'add')


def stop(ignore_errors=False):
    # Remove the tc classes
    popen([TC_CMD, 'qdisc', 'del', 'dev', os.environ['INTERFACE'], 'root'], ignore_errors=ignore_errors)


def update_rates(config):
    apply_rules(config[0], config[1], 'replace')


def usage():
    sys.exit("Usage: %s {start|stop|update-rates}" % sys.argv[0])


if __name__ == "__main__":
    if len(sys.argv) != 2:
        usage()

    if sys.argv[1] == "start":
        start(get_config_data(os.environ['ID']))
    elif sys.argv[1] == "stop":
        stop()
    elif sys.argv[1] == "update-rates":
        update_rates(get_config_data(os.environ['ID']))
    else:
        usage()
