1 files changed, 59 insertions, 0 deletions
diff --git a/src/home/mqtt/_wrapper.py b/src/home/mqtt/_wrapper.py
new file mode 100644
index 0000000..f858f88
--- /dev/null
+++ b/src/home/mqtt/_wrapper.py
@@ -0,0 +1,59 @@
+import paho.mqtt.client as mqtt
+
+from ._mqtt import Mqtt
+from ._node import MqttNode
+from ..config import config
+from ..util import strgen
+
+
+class MqttWrapper(Mqtt):
+    _nodes: list[MqttNode]
+
+    def __init__(self,
+                 client_id: str,
+                 topic_prefix='hk',
+                 randomize_client_id=False,
+                 clean_session=True):
+        if randomize_client_id:
+            client_id += '_'+strgen(6)
+        super().__init__(clean_session=clean_session,
+                         client_id=client_id)
+        self._nodes = []
+        self._topic_prefix = topic_prefix
+
+    def on_connect(self, client: mqtt.Client, userdata, flags, rc):
+        super().on_connect(client, userdata, flags, rc)
+        for node in self._nodes:
+            node.on_connect(self)
+
+    def on_disconnect(self, client: mqtt.Client, userdata, rc):
+        super().on_disconnect(client, userdata, rc)
+        for node in self._nodes:
+            node.on_disconnect()
+
+    def on_message(self, client: mqtt.Client, userdata, msg):
+        try:
+            topic = msg.topic
+            for node in self._nodes:
+                node.on_message(topic[len(f'{self._topic_prefix}/{node.id}/'):], msg.payload)
+        except Exception as e:
+            self._logger.exception(str(e))
+
+    def add_node(self, node: MqttNode):
+        self._nodes.append(node)
+        if self._connected:
+            node.on_connect(self)
+
+    def subscribe(self,
+                  node_id: str,
+                  topic: str,
+                  qos: int):
+        self._client.subscribe(f'{self._topic_prefix}/{node_id}/{topic}', qos)
+
+    def publish(self,
+                node_id: str,
+                topic: str,
+                payload: bytes,
+                qos: int):
+        self._client.publish(f'{self._topic_prefix}/{node_id}/{topic}', payload, qos)
+        self._client.loop_write()