diff --git a/collector.py b/collector.py index 1ca36d3..43a4c63 100644 --- a/collector.py +++ b/collector.py @@ -15,6 +15,7 @@ from plugins.system.memory_plugin import MemoryPlugin from plugins.system.network_plugin import NetworkPlugin from plugins.system.ping_plugin import PingPlugin from plugins.system.temperatures_plugin import TemperaturesPlugin +from plugins.system.speedtest_plugin import SpeedtestPlugin class Collector(object): @@ -29,6 +30,7 @@ class Collector(object): NetworkPlugin(), TemperaturesPlugin(), PingPlugin(), + SpeedtestPlugin(), # finance StocksPlugin(), @@ -46,9 +48,10 @@ class Collector(object): def schedule_plugins(self): start_date = datetime.now() + timedelta(seconds=10) for plugin in self.plugins: - self.scheduler.add_job(plugin.execute, 'interval', + self.scheduler.add_job(plugin.execute_wrapper, 'interval', seconds=plugin.get_interval(), - start_date=start_date) + start_date=start_date, + name=plugin.__class__.__name__) def run(self): logging.basicConfig() @@ -56,7 +59,8 @@ class Collector(object): models = self.collect_models() database.initialize_db() - database.DB.create_tables(models) + with database.DB.connection_context(): + database.DB.create_tables(models) self.schedule_plugins() logging.info('Started.') diff --git a/config.py b/config.py index 39cad52..a7a699d 100644 --- a/config.py +++ b/config.py @@ -59,6 +59,9 @@ PING_HOSTS = [ 'tibich.com' ] +### Speedtest +SPEEDTEST_INTERVAL = 15 * 60 # every 15 min + ### Stocks STOCKS_INTERVAL = 12 * 60 * 60 # updates daily diff --git a/database.py b/database.py index de9b74a..5a1e826 100644 --- a/database.py +++ b/database.py @@ -1,5 +1,6 @@ from peewee import DatabaseProxy, Model from playhouse.db_url import connect +from playhouse.pool import PooledPostgresqlExtDatabase import config DB = DatabaseProxy() diff --git a/install.sh b/install.sh index 88878bb..49c9d02 100755 --- a/install.sh +++ b/install.sh @@ -33,6 +33,7 @@ After=network.target [Service] Type=simple +User=$USER Restart=always RestartSec=1 ExecStart=/usr/bin/env python3 $(pwd)/collector.py diff --git a/plugins/plugin.py b/plugins/plugin.py index 61035c8..64931b6 100644 --- a/plugins/plugin.py +++ b/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() \ No newline at end of file diff --git a/plugins/system/smart_plugin.py b/plugins/system/smart_plugin.py new file mode 100644 index 0000000..b5e4b13 --- /dev/null +++ b/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() diff --git a/plugins/system/speedtest_plugin.py b/plugins/system/speedtest_plugin.py new file mode 100644 index 0000000..466968a --- /dev/null +++ b/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() diff --git a/requirements.txt b/requirements.txt index ccc9ec1..d4cbe34 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,9 +7,12 @@ psycopg2-binary # Used in most system plugins psutil +# Used by SMART plugin, requires 'smartmontools' installed + root access +pySMART + # robor plugin requests lxml # stocks plugin -yfinance \ No newline at end of file +yfinance