Update system api

rpi/services/mqtt_service.py

@@ -1,32 +1,96 @@
+from typing import Callable, Optional
+
 import paho.mqtt.client as mqtt
 import json
 import time
 import base64
 
+
 class MQTTService:
 
-    def __init__(self, address: str, port: int):
+    client : mqtt.Client
+
+    def __init__(self,
+                 address: str, port: int, client_id: Optional[str] = None,
+                 username: Optional[str] = None, password: Optional[str] = None,
+                 ca_certs: Optional[str] = None, insecure : bool = False
+                 ):
         self.address = address
         self.port = port
+        self.client = mqtt.Client(client_id=client_id)
 
-    def publish(self, client_id: str, topic: str, data: str, qos: int = 0):
-        client = mqtt.Client(client_id=client_id)
+        if username and password:
+            self.client.username_pw_set(username, password)
+
+        if ca_certs:
+            self.client.tls_set(ca_certs=ca_certs)
+            self.client.tls_set()
+            self.client.tls_insecure_set(insecure)
+
+        self._connected = False
+        self.client.on_connect = self._on_connect
+        self.client.on_disconnect = self._on_disconnect
+        self.client.on_message = None
+        self._subscriptions = {}
+
+    def _on_connect(self, client, userdata, flags, rc):
+        if rc == 0:
+            self._connected = True
+            for topic in self._subscriptions:
+                self.client.subscribe(topic)
+        else:
+            print(f"MQTT connection failed with code {rc}")
+
+    def _on_disconnect(self, client, userdata, rc):
+        self._connected = False
+
+    def connect(self):
+        if not self._connected:
+            print(f"Connecting to {self.address}...")
+            self.client.connect(self.address, self.port)
+            self.client.loop_start()
+            timeout = 5
+            start = time.time()
+            while not self._connected and time.time() - start < timeout:
+                time.sleep(0.1)
+            if not self._connected:
+                raise ConnectionError(f"Cannot connect to MQTT broker at {self.address}:{self.port}")
+            print(f"Successfully connected to {self.address}")
+        else :
+            print(f"Already connected to {self.address}...")
+
+    def disconnect(self):
+        if self._connected:
+            self.client.disconnect()
+            self.client.loop_stop()
+            self._connected = False
+
+    def publish(self, topic: str, data: str, qos: int = 0):
         try:
-            client.connect(self.address, self.port)
-            client.loop_start()
-
+            self.connect()
             payload = {
                 "timestamp": int(time.time()),
                 "data": data
             }
-            result = client.publish(topic, json.dumps(payload), qos=qos)
+            result = self.client.publish(topic, json.dumps(payload), qos=qos)
             result.wait_for_publish()
-
         except Exception as e:
-            print("Erreur MQTT:", e)
-        finally:
-            client.loop_stop()
-            client.disconnect()
+            print(f"An error occurred while publishing on {topic} on {self.address} : {e}")
+
+    def subscribe(self, topic: str, handler: Callable[[str, str], None], qos: int = 0):
+        try:
+            self.connect()
+            self._subscriptions[topic] = handler
+
+            def on_message(client, userdata, msg):
+                for sub_topic, h in self._subscriptions.items():
+                    if mqtt.topic_matches_sub(sub_topic, msg.topic):
+                        h(msg.topic, msg.payload.decode())
+
+            self.client.on_message = on_message
+            self.client.subscribe(topic, qos=qos)
+        except Exception as e:
+            print(f"An error occurred while trying to subscribe to {topic} on {self.address} : {e}")
 
 
 if __name__ == "__main__":