Added speedtest plugin, fixed database connection error.

plugins/plugin.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from typing import Tuple
+import database
 
 
 class Plugin(ABC):
@@ -12,3 +13,7 @@ class Plugin(ABC):
     @abstractmethod
     def execute(self) -> None:
         pass
+
+    def execute_wrapper(self) -> None:
+        with database.DB.connection_context():
+            self.execute()

plugins/system/smart_plugin.py

@@ -0,0 +1,39 @@
+from datetime import datetime
+
+import psutil
+from peewee import *
+from playhouse.shortcuts import model_to_dict
+
+import config
+from database import BaseModel
+from plugins.plugin import Plugin
+
+
+class SMART(BaseModel):
+    time = DateTimeField(index=True, default=datetime.utcnow)
+    drive = TextField(null=False)
+    attribute_id = IntegerField(null=False)
+    attribute_name = TextField(null=False)
+    value = IntegerField(null=False)
+    worst = IntegerField(null=False)
+    threshold = IntegerField(null=False)
+    raw = IntegerField(null=False)
+
+
+class SMARTPlugin(Plugin):
+    models = [SMART]
+
+    def get_interval(self):
+        return config.DISK_USAGE_INTERVAL
+
+    def execute(self):
+        for partition in psutil.disk_partitions():
+            usage = psutil.disk_usage(partition.mountpoint)
+            
+            entry = DiskUsage()
+            entry.partition = partition.device
+            entry.mountpoint = partition.mountpoint
+            entry.total = usage.total
+            entry.used = usage.used
+            entry.free = usage.free
+            entry.save()

plugins/system/speedtest_plugin.py

@@ -0,0 +1,57 @@
+import subprocess
+import re
+import subprocess
+from datetime import datetime
+
+import psutil
+from peewee import *
+from playhouse.shortcuts import model_to_dict
+
+import config
+from database import BaseModel
+from plugins.plugin import Plugin
+
+import json
+
+import logging
+
+class Speedtest(BaseModel):
+    time = DateTimeField(index=True, default=datetime.utcnow)
+    upload = IntegerField(null=False)
+    download = IntegerField(null=False)
+    latency = FloatField(null=True)
+    jitter = FloatField(null=True)
+    packetLoss = IntegerField(null=True)
+
+
+class SpeedtestPlugin(Plugin):
+    models = [Speedtest]
+
+    def get_interval(self):
+        return config.SPEEDTEST_INTERVAL
+
+    def execute(self):
+        command = ['/usr/local/bin/SpeedTest', '--output', 'json']
+        proc = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        stdout = proc.stdout.decode()
+        stderr = proc.stderr.decode()
+
+        entry = Speedtest()
+        entry.download = 0
+        entry.upload = 0
+
+        if proc.returncode == 0:
+            try:
+                result = json.loads(stdout)
+                entry.download = int(float(result['download']))
+                entry.upload = int(float(result['upload']))
+                entry.latency = float(result['ping'])
+                entry.jitter = float(result['jitter'])
+            except BaseException as e:
+                logging.error(f"SpeedTest failed: {e}")
+
+        else:
+            logging.error(f"SpeedTest nonzero return: {proc.returncode}\n-----\n{stdout}\n{stderr}\n\n")
+            
+
+        entry.save()