BroSpeedTester/speedtester.py

from configuration import read_config
# region Logger
import logging
from debug import setup_logging, catch_errors

log = logger = logging.getLogger("default")
setup_logging()
# endregion

from dbo import Entry
import routeros_api

config = read_config()
secrets = read_config('secrets')

import time


@catch_errors
def gather_data():
    log.debug("Gathering data...")
    downloads = []
    uploads = []
    dates = []
    for entry in Entry.select():
        downloads.append(entry.download)
        uploads.append(entry.upload)
        dates.append(entry.date_created)

    return dates, downloads, uploads


@catch_errors
def generate_plot_image(dates, downloads, uploads):
    log.debug("Genering image output...")
    import matplotlib
    import matplotlib.pyplot as plt

    dates = matplotlib.dates.date2num(dates)
    fig = plt.figure(figsize=(12, 3))
    plt.plot_date(dates, downloads, fmt="b-")
    plt.ylabel('Download Speed Mbps')
    plt.tight_layout()
    plt.savefig(read_config()['output_image_path'])


@catch_errors
def generate_txt_output(dates, downloads, uploads):
    log.debug("Genering txt output...")
    txt = "Date:			Down;			Up;\n"
    for i, date in enumerate(dates):
        download = downloads[i]
        upload = uploads[i]

        txt += f"{date}:			{download} Mbps;			{upload} Mbps\n"
    with open(read_config()['output_txt_path'], "w+") as f:
        f.write(txt)


@catch_errors
def ros_fastrack_enable(enable):
    '''
    Enable or disable fasttrack
    :param enable: True to enable, False to disable
    '''

    if enable:
        disabled = 'false'
    else:
        disabled = 'true'

    log.info(f"Changing fasttrack disabled to {disabled}")

    connection = routeros_api.RouterOsApiPool(config['ros_ip'], username=secrets["ros_login"],
                                              password=secrets["ros_password"], plaintext_login=True)
    api = connection.get_api()

    filter = api.get_resource('/ip/firewall/filter')

    for i in filter.get():
        if config['ros_fasttrack_comment'] in i['comment']:
            filter.set(id=i['id'], disabled=disabled)
    connection.disconnect()


@catch_errors
def ros_dynamic_speed(upload, download):
    '''
    Adjust router os queue speed.
    :param upload: Upload speed from speedtest
    :param download: Download speed from speedtest
    '''

    connection = routeros_api.RouterOsApiPool(config['ros_ip'], username=secrets["ros_login"],
                                              password=secrets["ros_password"], plaintext_login=True)
    api = connection.get_api()

    list_queues = api.get_resource('/queue/simple')

    for queue in list_queues.get():
        if queue['name'] in config['ros_queues']:
            log.debug(
                f"Adjust Queue {queue['name']}: limit_at {int(upload) - int(upload / 10)}/{int(download) - int(download / 10)}; max_limit {int(upload)}/{int(download)}")

            if config["ros_du_invert"] == True:
                # Inverting upload and download values, because master queue is most likely applied to the bridge
                list_queues.set(id=queue['id'], max_limit=f"{int(download)}/{int(upload)}")
            else:
                # Not inverting, use this in case master queue is applied to something like LTE interface
                list_queues.set(id=queue['id'], max_limit=f"{int(upload)}/{int(download)}")
    connection.disconnect()


def speedtest():
    import speedtest

    servers = []
    threads = None

    for i in range(0, 3):
        try:
            log.debug("Initializing speedtest...")
            s = speedtest.Speedtest()

            log.debug(f"Running test...")
            s.get_servers(servers)
            s.get_best_server()
            s.download(threads=threads)
            s.upload(threads=threads, pre_allocate=False)

            results_dict = s.results.dict()

            if results_dict['download'] >= config['ros_minimum_speed']:
                download = round(results_dict['download'] / 1000000, 2)
                upload = round(results_dict['upload'] / 1000000, 2)
                break
            time.sleep(10)
        except:
            log.error(f"Test failed, try {i + 1}/3", exc_info=True)

    log.debug(f"{download}mbps, {upload}mbps")
    return download, upload, results_dict

if __name__ == "__main__":
    '''
    This script will run a few speed tests, calculate average upload and download speeds and record them into database.
    Once finished it will also generate an image with graph plotted.
    '''

    # generate_plot_image(*gather_data())
    # import sys
    # sys.exit()
    from random import uniform

    try:
        if config["ros_dynamic_speed"]:
            ros_fastrack_enable(True)
            time.sleep(5)

        download, upload, results_dict = speedtest()

        entry = Entry()
        entry.upload = upload
        entry.download = download
        entry.save()

        ros_upload = results_dict['upload']
        ros_download = results_dict['download']

        if ros_download < config['ros_minimum_speed']:
            ros_download = config['ros_minimum_speed']

        if config["ros_dynamic_speed"]:
            ros_dynamic_speed(results_dict['upload'], results_dict['download'])
            ros_fastrack_enable(False)

        dates, downloads, uploads = gather_data()

        generate_txt_output(dates, downloads, uploads)

        generate_plot_image(dates, downloads, uploads)


    except:
        log.error("Error!", exc_info=True)

        if config["ros_dynamic_speed"]:
            ros_fastrack_enable(False)