Add dynamic speed test and queue adjustment

.gitignore

@@ -131,3 +131,6 @@ dmypy.json
 
 .idea/
 data/database.db
+secrets.yaml
+test.py
+test.png

config/config.yaml

@@ -1,3 +1,9 @@
 database_migrate: False
-output_image_path: /var/www/downloads/speedgraph.png
+output_image_path: speedgraph.png
-output_txt_path: /var/www/downloads/speeds.txt
+output_txt_path: speeds.txt
+ros_dynamic_speed: True
+ros_ip: 192.168.88.1
+ros_queues:
+  - WAN
+ros_du_invert: False
+ros_fasttrack_comment: "defconf: fasttrack"

debug.py

@@ -40,18 +40,6 @@ def setup_logging(
         logging.basicConfig(level=default_level)
 
 
-def catch_errors_json(f):
-    @functools.wraps(f)
-    def wrapped(*args, **kwargs):
-        try:
-            return f(*args, **kwargs)
-        except Exception as e:
-            traceback.print_exc()
-            return jsonify({"error": str(e), "traceback": traceback.format_exc()})
-
-    return wrapped
-
-
 loggers = {}
 
 
@@ -70,7 +58,7 @@ def get_logger(name):
 
 log = logger = get_logger("default")
 
-def catch_errors_json(f):
+def catch_errors(f):
     @functools.wraps(f)
     def wrapped(*args, **kwargs):
         try:
@@ -78,18 +66,6 @@ def catch_errors_json(f):
         except Exception as e:
             traceback.print_exc()
             log.error(traceback.format_exc())
-            return jsonify({"error": str(e), "traceback": traceback.format_exc()})
+            return traceback.format_exc()
 
-    return wrapped
+    return wrapped
-
-def catch_errors_html(f):
-    @functools.wraps(f)
-    def wrapped(*args, **kwargs):
-        try:
-            return f(*args, **kwargs)
-        except Exception as e:
-            traceback.print_exc()
-            log.error(traceback.format_exc())
-            return render_template("error.html", error=str(e), error_trace=traceback.format_exc())
-
-    return wrapped

requirements.txt

@@ -1,4 +1,5 @@
 peewee
 speedtest-cli
 numpy
-matplotlib
+matplotlib
+RouterOS-api

speeds.txt

@@ -0,0 +1,4 @@
+Date:			Down;			Up;
+2020-06-24 02:37:11.363969:			6.29 Mbps;			4.16 Mbps
+2020-06-24 02:38:12.903694:			5.52 Mbps;			3.8 Mbps
+2020-06-24 02:39:31.212599:			12.07 Mbps;			5.08 Mbps

speedtester.py

@@ -1,15 +1,22 @@
 from configuration import read_config
 # region Logger
 import logging
-from debug import setup_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 = []
@@ -22,18 +29,22 @@ def gather_data():
 
     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))
+    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"
@@ -45,6 +56,61 @@ def generate_txt_output(dates, downloads, uploads):
     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()
+
+
 if __name__ == "__main__":
     '''
     This script will run a few speed tests, calculate average upload and download speeds and record them into database.
@@ -55,18 +121,22 @@ if __name__ == "__main__":
     # import sys
     # sys.exit()
     from random import uniform
+
     try:
+        if config["ros_dynamic_speed"]:
+            ros_fastrack_enable(True)
+            time.sleep(5)
+
         import speedtest
 
         servers = []
         threads = None
 
-        for i in range (0, 3):
+        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()
@@ -74,13 +144,13 @@ if __name__ == "__main__":
                 s.upload(threads=threads, pre_allocate=False)
 
                 results_dict = s.results.dict()
-                download = round(results_dict['download']/1000000, 2)
+                download = round(results_dict['download'] / 1000000, 2)
-                upload = round(results_dict['upload']/1000000, 2)
+                upload = round(results_dict['upload'] / 1000000, 2)
                 # download = uniform(0,2)
                 # upload = uniform(0,2)
                 break
             except:
-                log.error(f"Test failed, try {i+1}/3", exc_info=True)
+                log.error(f"Test failed, try {i + 1}/3", exc_info=True)
 
         log.debug(f"{download}mbps, {upload}mbps")
 
@@ -88,22 +158,21 @@ if __name__ == "__main__":
         entry.upload = upload
         entry.download = download
         entry.save()
-    except:
-        log.error("Data record error.", exc_info=True)
 
-    try:
+        if config["ros_dynamic_speed"]:
+            ros_dynamic_speed(results_dict['upload'], results_dict['download'])
+
+            ros_fastrack_enable(False)
+
         dates, downloads, uploads = gather_data()
 
-        try:
+        generate_txt_output(dates, downloads, uploads)
-            generate_txt_output(dates, downloads, uploads)
+
-        except:
+        generate_plot_image(dates, downloads, uploads)
-            log.error("Unable to save text file.", exc_info=True)
 
-        try:
-            generate_plot_image(dates, downloads, uploads)
-        except:
-            log.error("Unable to save plot file.", exc_info=True)
 
     except:
-        log.error("Error plotting.", exc_info=True)
+        log.error("Error!", exc_info=True)
 
+        if config["ros_dynamic_speed"]:
+            ros_fastrack_enable(False)