BroSpeedTester/speedtester.py

374 lines
11 KiB
Python

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 statistics import median
from threading import Thread
import speedtest
from dbo import Entry
import routeros_api
config = read_config()
secrets = read_config('secrets')
import time
def mbits(bits):
return round(bits / 1000000, 2)
@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
import matplotlib
import matplotlib.pyplot as plt
@catch_errors
def generate_plot_image(dates, downloads, uploads, name="speed", description="Speed Graph"):
log.debug("Genering image output for {}...".format(name))
dates = matplotlib.dates.date2num(dates)
fig = plt.figure(figsize=(len(dates / 10), 3))
plt.plot_date(dates, downloads, fmt="b-")
plt.ylabel(description + " (Mbps)")
plt.tight_layout()
plt.savefig(read_config()['output_image_path'].format(name))
@catch_errors
def generate_diff(dates, downloads, uploads):
dl = []
up = []
for i, date in enumerate(dates):
curr = downloads[i] if downloads[i] else 0
prev = downloads[i-1] if downloads[i-1] and i > 0 else 0
dl.append(curr - prev)
curr = downloads[i] if downloads[i] else 0
prev = downloads[i - 1] if downloads[i - 1] and i > 0 else 0
up.append(curr - prev)
generate_plot_image(dates, dl, up, "diff", "Difference Graph")
def generate_updown_plot_simple(downloads, uploads, name, description):
fig = plt.figure(figsize=(10, 3))
plt.plot(downloads, linestyle='dashed', color="red")
plt.plot(uploads, linestyle='dashed', color="green")
plt.ylabel(description + " (Mbps)")
plt.tight_layout()
plt.savefig(read_config()['output_image_path'].format(name))
@catch_errors
def gather_week_median_data(dates, downloads, uploads):
log.debug("Gather week median data...")
dl = []
up = []
for hour in range(24):
dl.append([])
up.append([])
for i, date in enumerate(dates):
hour = date.hour
dl[hour].append(downloads[i] if downloads[i] else 0)
up[hour].append(uploads[i] if downloads[i] else 0)
for i, hour in enumerate(dl):
dl[i] = median(dl[i]) if len(dl[i]) > 0 else 0
for i, hour in enumerate(up):
up[i] = median(up[i]) if len(up[i]) > 0 else 0
return dl, up
@catch_errors
def generate_week_median(dates, downloads, uploads):
dl, up = gather_week_median_data(dates, downloads, uploads)
generate_updown_plot_simple(dl, up, "week_median", "Week median")
@catch_errors
def generate_week_median_diff(dates, downloads, uploads):
dl, up = gather_week_median_data(dates, downloads, uploads)
downs = []
ups = []
for i, down in enumerate(dl):
if i > 0:
downs.append(dl[i]-dl[i-1])
ups.append(up[i]-up[i-1])
else:
downs.append(down)
ups.append(up[i])
generate_updown_plot_simple(downs, ups, "week_median_diff", "Week median diff")
@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
'''
log.debug(f"Set Dynamic Speed to: DOWN {mbits(download)} mbps; UP {mbits(upload)} mbps")
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']}: 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()
wan_upload = None
wan_download = None
results_dict = None
test_started = False
speedtest_failed = False
downloading = True # True for download, False for upload
threads = None
servers = []
def reset_globals():
global wan_upload
global wan_download
global results_dict
global test_started
global speedtest_failed
global downloading
wan_upload = None
wan_download = None
results_dict = None
test_started = False
speedtest_failed = False
downloading = True
def threaded_speedtest():
global test_started
global downloading
global results_dict
global speedtest_failed
try:
s = speedtest.Speedtest()
s.get_servers(servers)
s.get_best_server()
print(f"Running test...")
test_started = True
downloading = True
s.download(threads=threads)
downloading = False
s.upload(threads=threads, pre_allocate=False)
results_dict = s.results.dict()
log.info(
f"Speedtest.net result: DOWN {mbits(results_dict['download'])} mbps; UP {mbits(results_dict['upload'])} mbps;")
except Exception as e:
speedtest_failed = True
log.error(f"Speedtest.net failed: {e}", exc_info=True)
return results_dict
def threaded_wan_speed():
global test_started
global results_dict
global downloading
print("Waiting for test to start...")
while not test_started:
if speedtest_failed:
return
time.sleep(1)
print("Allow warm-up...")
time.sleep(2)
uploads = []
downloads = []
print("Monitoring...")
while results_dict == None:
connection = routeros_api.RouterOsApiPool(config['ros_ip'], username=secrets["ros_login"],
password=secrets["ros_password"], plaintext_login=True)
api = connection.get_api()
traffic = api.get_resource('/').call('interface/monitor-traffic',
{'interface': config['ros_wan_interface'], 'once': ' '})[0]
traffic['rx-bits-per-second'] = int(traffic['rx-bits-per-second'])
traffic['tx-bits-per-second'] = int(traffic['tx-bits-per-second'])
if downloading:
downloads.append(traffic['rx-bits-per-second'])
print(f"DL: {mbits(traffic['rx-bits-per-second'])} mbps;")
else:
uploads.append(traffic['tx-bits-per-second'])
print(f"UP: {mbits(traffic['tx-bits-per-second'])} mbps;")
time.sleep(1)
global wan_download
global wan_upload
wan_download = median(downloads)
wan_upload = median(uploads)
log.info(f"Monitor result: {mbits(wan_download)} mbps; {mbits(wan_upload)} mbps;")
def test_speed():
global wan_download
tries = 2
for i in range(tries):
reset_globals()
sws = Thread(target=threaded_wan_speed)
st = Thread(target=threaded_speedtest)
st.start()
sws.start()
st.join()
sws.join()
if wan_download >= config['ros_doublecheck_speed']:
break
elif i < tries - 1:
log.warning("Speed is below ros_doublecheck_speed limit, retrying in a few seconds.")
time.sleep(10)
return
def generate_database_reports():
dates, downloads, uploads = gather_data()
generate_txt_output(dates, downloads, uploads)
generate_plot_image(dates, downloads, uploads)
generate_week_median(dates, downloads, uploads)
generate_week_median_diff(dates, downloads, uploads)
generate_diff(dates, downloads, uploads)
import sys
from netutils import test_intertnet_connection
def on_fail_or_no_connection():
ros_fastrack_enable(False)
entry = Entry()
entry.upload = None
entry.download = None
entry.save()
generate_database_reports()
log.warning("No internet connection! Exiting.")
sys.exit()
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:
log.debug("Test internet connection...")
if config["ros_dynamic_speed"]:
ros_fastrack_enable(True)
time.sleep(5)
test_speed()
if speedtest_failed:
on_fail_or_no_connection()
entry = Entry()
entry.upload = mbits(wan_upload)
entry.download = mbits(wan_download)
entry.save()
if wan_download < config['ros_minimum_speed']:
wan_download = config['ros_minimum_speed']
if config["ros_dynamic_speed"]:
ros_dynamic_speed(wan_upload, wan_download)
ros_fastrack_enable(False)
generate_database_reports()
except:
log.error("Error!", exc_info=True)
if config["ros_dynamic_speed"]:
ros_fastrack_enable(False)