#!/usr/bin/python3

import sys
import os.path
import time
import logging
import argparse
import serial
import subprocess
import re

def get_bytes(t, iface='enp0s25'):
    with open('/sys/class/net/' + iface + '/statistics/' + t + '_bytes', 'r') as f:
        data = f.read();
    return int(data)

def get_link_speed(iface='enp0s25'):
    ethtool=subprocess.Popen(['ethtool',iface], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    s = re.search('Speed: (?P<Speed>\d+)Mb/s', ethtool.stdout.read().decode('utf-8'))
    return int(s.group('Speed')) * 125000

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-v', '--verbose', action='count', default=0,
                        help='increase logging verbosity')
    parser.add_argument('-s', '--serial', type=str, default='/dev/ttyACM0',
                        help='serial port for arduino connection')
    parser.add_argument('-t', '--time', type=float, default=0.02,
                        help='update interval for measurements')
    parser.add_argument('interfaces', metavar='interface', type=str, nargs='+',
                        help='network interfaces to take into account')
    args = parser.parse_args()

    # Process verbosity
    if args.verbose == 0:
        LOGLEVEL = logging.WARNING
    elif args.verbose == 1:
        LOGLEVEL = logging.INFO
    elif args.verbose >= 2:
        LOGLEVEL = logging.DEBUG

    logging.basicConfig(format='%(levelname)s:%(name)s:%(asctime)s:%(message)s', level=LOGLEVEL)
    logging.info("Network utilization visualization starting up")

    if not os.path.exists(args.serial):
        logging.error("Serial output not available, quitting!")
        sys.exit(1)
    else:
        logging.info("Opening serial port %s", args.serial)
        port = serial.Serial(
            port=args.serial,
            baudrate=115200,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            timeout=1,
            xonxoff=False,
            rtscts=False,
            dsrdtr=False,
            write_timeout=None,
            inter_byte_timeout=None)

        time.sleep(2)

        logging.info("Serial port opened")

        sum_max = 0
        ifaces = []
        for iface in args.interfaces:
            ifaces.append({
                'name': iface,
                'rx': get_bytes('rx', iface),
                'tx': get_bytes('tx', iface)})
            sum_max += get_link_speed(iface) * args.time

        while True:
            sum_rx = 0;
            sum_tx = 0;

            for iface in ifaces:
                if_rx = get_bytes('rx', iface['name'])
                if_tx = get_bytes('tx', iface['name'])
                sum_rx += if_rx - iface['rx']
                sum_tx += if_tx - iface['tx']
                iface['rx'] = if_rx
                iface['tx'] = if_tx

            message = bytearray(3)
            message[0] = 0xFF
            message[1] = int(sum_tx * 100 / sum_max)
            message[2] = int(sum_rx * 100 / sum_max)

            port.write(message)

            time.sleep(args.time)