Smartwave/shared/mqtt.py

"""Small MQTT compatibility layer for CPython and MicroPython.

The wrapper keeps the broker host explicit so embedded clients can point to a
real IP address instead of localhost.
"""

try:
    import paho.mqtt.client as _mqtt
    BACKEND_NAME = "paho"
    IS_MICROPYTHON = False
except ImportError:
    try:
        from umqtt.simple import MQTTClient as _MQTTClient
        BACKEND_NAME = "umqtt.simple"
        IS_MICROPYTHON = True
    except ImportError:
        try:
            from umqtt.robust import MQTTClient as _MQTTClient
            BACKEND_NAME = "umqtt.robust"
            IS_MICROPYTHON = True
        except ImportError as exc:
            raise ImportError("No MQTT client found. Expected paho.mqtt or umqtt.") from exc


DEFAULT_PORT = 8884
DEFAULT_CA_FILE = "orchestrateur/mqtt/certs/ca.crt"


def _resolve_host(host):
    if host:
        return host

    try:
        import os

        getenv = getattr(os, "getenv", None)
        if getenv is not None:
            host = getenv("MQTT_BROKER_HOST")
    except Exception:
        host = None

    if not host:
        raise ValueError("MQTT broker host is required. Pass the broker IP address instead of localhost.")

    return host


def _ensure_bytes(payload):
    if payload is None:
        return b""
    if isinstance(payload, bytes):
        return payload
    if isinstance(payload, bytearray):
        return bytes(payload)
    return str(payload).encode()


def _read_file_bytes(path):
    with open(path, "rb") as handle:
        return handle.read()


class BrokerClient:
    """Small MQTT client with a normalised API across runtimes."""

    def __init__(self, host=None, port=DEFAULT_PORT, client_id=None, use_tls=True, cafile=None, certfile=None, keyfile=None, ssl_params=None, tls_insecure=False, username=None, password=None, keepalive=60):
        self.host = _resolve_host(host)
        self.port = port
        self.client_id = client_id
        self.use_tls = use_tls
        self.cafile = cafile or DEFAULT_CA_FILE
        self.certfile = certfile
        self.keyfile = keyfile
        self.ssl_params = ssl_params
        self.tls_insecure = tls_insecure
        self.username = username
        self.password = password
        self.keepalive = keepalive
        self._client = None
        self._callback = None
        self._messages = []

    def set_callback(self, callback):
        self._callback = callback
        if self._client is not None and not IS_MICROPYTHON:
            self._client.on_message = self._on_message

    def _store_message(self, topic, payload, qos=None, retain=False):
        message = {
            "topic": topic,
            "payload": payload,
            "qos": qos,
            "retain": retain,
        }
        self._messages.append(message)
        if self._callback is not None:
            self._callback(message)

    def _on_message(self, client, userdata, msg):
        self._store_message(msg.topic, msg.payload, getattr(msg, "qos", None), getattr(msg, "retain", False))

    def open(self):
        if self._client is not None:
            return self._client

        if IS_MICROPYTHON:
            ssl_params = self.ssl_params
            if self.use_tls and ssl_params is None and self.cafile is not None:
                ssl_params = {"cadata": _read_file_bytes(self.cafile)}
            client = _MQTTClient(
                self.client_id or "smartWave-client",
                self.host,
                port=self.port,
                user=self.username,
                password=self.password,
                keepalive=self.keepalive,
                ssl=self.use_tls or ssl_params is not None,
                ssl_params=ssl_params,
            )
            self._client = client
            return self._client

        client = _mqtt.Client(client_id=self.client_id or "", clean_session=True, protocol=4, transport="tcp")
        if self.username is not None or self.password is not None:
            client.username_pw_set(self.username, self.password)
        if self.use_tls:
            tls_kwargs = {}
            if self.cafile is not None:
                tls_kwargs["ca_certs"] = self.cafile
            if self.certfile is not None:
                tls_kwargs["certfile"] = self.certfile
            if self.keyfile is not None:
                tls_kwargs["keyfile"] = self.keyfile
            if tls_kwargs:
                client.tls_set(**tls_kwargs)
            else:
                client.tls_set()
            if self.tls_insecure:
                client.tls_insecure_set(True)
        client.on_message = self._on_message
        self._client = client
        return self._client

    def connect(self):
        client = self.open()
        if IS_MICROPYTHON:
            client.connect()
            return client

        client.connect(self.host, self.port, self.keepalive)
        return client

    def publish(self, topic, payload, qos=2, retain=False):
        client = self.open()
        payload_bytes = _ensure_bytes(payload)
        if IS_MICROPYTHON:
            return client.publish(topic, payload_bytes, retain=retain, qos=qos)
        return client.publish(topic, payload_bytes, qos=qos, retain=retain)

    def subscribe(self, topic, qos=2):
        client = self.open()
        if IS_MICROPYTHON:
            client.set_callback(self._on_micropython_message)
            return client.subscribe(topic, qos=qos)
        return client.subscribe(topic, qos=qos)

    def _on_micropython_message(self, topic, payload):
        self._store_message(topic, payload, None, False)

    def poll(self, timeout=0.1):
        if self._client is None:
            return None
        if IS_MICROPYTHON:
            return self._client.check_msg()
        return self._client.loop(timeout=timeout)

    def wait(self):
        if self._client is None:
            return None
        if IS_MICROPYTHON:
            return self._client.wait_msg()
        return self._client.loop_forever()

    def get_message(self):
        if not self._messages:
            return None
        return self._messages.pop(0)

    def close(self):
        if self._client is None:
            return
        try:
            self._client.disconnect()
        except Exception:
            pass
        self._client = None

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc, traceback):
        self.close()


def connect(host=None, **client_kwargs):
    return BrokerClient(host=host, **client_kwargs)


def publish(host, topic, payload, **client_kwargs):
    qos = client_kwargs.pop("qos", 2)
    retain = client_kwargs.pop("retain", False)
    client = connect(host=host, **client_kwargs)
    client.connect()
    try:
        return client.publish(topic, payload, qos=qos, retain=retain)
    finally:
        client.close()