Management command for the mqtt loop

8 files changed, 113 insertions, 0 deletions

diff --git a/kerbana/settings.py b/kerbana/settings.py
index 7be9578..f0ba3f2 100644
--- a/kerbana/settings.py
+++ b/kerbana/settings.py
@@ -37,6 +37,7 @@ INSTALLED_APPS = [
     'django.contrib.sessions',
     'django.contrib.messages',
     'django.contrib.staticfiles',
+    'rrd',
 ]
 
 MIDDLEWARE = [
@@ -123,3 +124,9 @@ STATIC_URL = '/static/'
 # https://docs.djangoproject.com/en/3.2/ref/settings/#default-auto-field
 
 DEFAULT_AUTO_FIELD = 'django.db.models.BigAutoField'
+
+MQTT_SERVER = "test.mosquitto.org"
+MQTT_PORT = 1883
+MQTT_USER = None
+MQTT_PASSWORD = None
+MQTT_TOPIC = "kerbana/#"
diff --git a/rrd/management/__init__.py b/rrd/management/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/rrd/management/__init__.py
diff --git a/rrd/management/commands/__init__.py b/rrd/management/commands/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/rrd/management/commands/__init__.py
diff --git a/rrd/management/commands/listen_mqtt.py b/rrd/management/commands/listen_mqtt.py
new file mode 100644
index 0000000..758a4f2
--- /dev/null
+++ b/rrd/management/commands/listen_mqtt.py
@@ -0,0 +1,10 @@
+from django.core.management.base import BaseCommand
+
+import rrd.mqtt
+
+class Command(BaseCommand):
+    help = "Listen to mqtt events and update rrd files and graphs"
+
+    def handle(self, *args, **kw):
+        mqtt_client = rrd.mqtt.MQTTClient()
+        mqtt_client.loop_forever()
diff --git a/rrd/mqtt.py b/rrd/mqtt.py
new file mode 100644
index 0000000..713d837
--- /dev/null
+++ b/rrd/mqtt.py
@@ -0,0 +1,64 @@
+import logging
+
+import django.conf
+import paho.mqtt.client as mqtt
+
+
+log = logging.getLogger(__name__)
+
+
+class MQTTClient:
+    def __init__(self):
+        self.reconnect = True
+        self.connected = False
+
+        self.client = mqtt.Client()
+        self.client.on_connect = self.on_connect
+        self.client.on_disconnect = self.on_disconnect
+        self.client.on_message = self.on_message
+
+    def loop_forever(self):
+        self.connect()
+        self.client.loop_forever()
+
+    def loop_start(self):
+        self.connect()
+        self.client.loop_start()
+
+    def loop_stop(self):
+        self.disconnect(reconnect=False)
+        self.client.loop_stop()
+
+    def connect(self):
+        try:
+            self.client.connect(
+                django.conf.settings.MQTT_SERVER,
+                django.conf.settings.MQTT_PORT,
+                60,  # TODO: make the keepalive configurable
+            )
+        except OSError as e:
+            log.debug("Could not connect to MQTT server")
+            log.debug(e)
+
+    def disconnect(self, reconnect: bool = True):
+        self.reconnect = reconnect
+        self.client.disconnect()
+
+    def on_connect(self, client, userdata, flags, rc):
+        log.info("Connected to MQTT")
+        self.connected = True
+        client.subscribe(django.conf.settings.MQTT_TOPIC)
+
+    def on_disconnect(self, client, userdata, rc):
+        log.info("Disconnected from MQTT")
+        self.connected = False
+        if self.reconnect:
+            self.connect()
+
+    def on_message(self, client, userdata, msg):
+        print(
+            "Received msg %s %s", msg.topic, msg.payload.decode()
+        )
+        log.debug(
+            "Received msg %s %s", msg.topic, msg.payload.decode()
+        )
diff --git a/rrd/tests/__init__.py b/rrd/tests/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/rrd/tests/__init__.py
diff --git a/rrd/tests.py b/rrd/tests/test_django.py
index 7ce503c..e0fbb45 100644
--- a/rrd/tests.py
+++ b/rrd/tests/test_django.py
@@ -1,3 +1,5 @@
 from django.test import TestCase
 
 # Create your tests here.
+
+
diff --git a/rrd/tests/test_mqtt.py b/rrd/tests/test_mqtt.py
new file mode 100644
index 0000000..4ea0852
--- /dev/null
+++ b/rrd/tests/test_mqtt.py
@@ -0,0 +1,30 @@
+import time
+
+import django.test
+
+from .. import mqtt
+
+class TestMQTT(django.test.TestCase):
+    def setUp(self):
+        self.mqtt = mqtt.MQTTClient()
+        self.mqtt.loop_start()
+        time.sleep(0.1)
+        if not self.mqtt.connected:
+            self.skipTest("Could not find an mqtt server")
+
+    def tearDown(self):
+        self.mqtt.loop_stop()
+
+    def test_disconnect(self):
+        # after disconnecting from the mqtt server, we should
+        # automatically reconnect
+        self.mqtt.disconnect()
+        time.sleep(2)
+        self.assertTrue(self.mqtt.connected)
+
+    def test_disconnect_and_stay(self):
+        # unless we really want to force a disconnection
+        self.mqtt.disconnect(reconnect=False)
+        time.sleep(2)
+        self.assertFalse(self.mqtt.connected)
+        self.assertFalse(self.mqtt.reconnect)