266 lines
8.4 KiB
Python
266 lines
8.4 KiB
Python
"""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:
|
|
import ssl
|
|
ssl_params = self.ssl_params
|
|
|
|
if self.use_tls and ssl_params is None:
|
|
# MicroPython uses context-less structures.
|
|
# If your CA is self-signed, validation can fail without a valid hostname match.
|
|
ssl_params = {
|
|
"cert_reqs": ssl.CERT_NONE, # Temporarily change to NONE to test if validation is the culprit
|
|
"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,
|
|
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)
|
|
|
|
if isinstance(topic, bytes):
|
|
topic = topic.decode('utf-8')
|
|
|
|
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)
|
|
|
|
if isinstance(topic, bytes):
|
|
topic = topic.decode('utf-8')
|
|
|
|
return client.subscribe(topic, qos=qos)
|
|
|
|
def unsubscribe(self, topic):
|
|
client = self.open()
|
|
if IS_MICROPYTHON:
|
|
import struct
|
|
# Ensure the topic is bytes for writing to the socket
|
|
topic_bytes = topic if isinstance(topic, bytes) else topic.encode('utf-8')
|
|
|
|
# 1. Build the MQTT unsubscribe packet header
|
|
pkt = bytearray(b"\xa2\0\0\0")
|
|
client.pid += 1
|
|
|
|
# Packet length is: 2 bytes (PID) + 2 bytes (topic length indicator) + topic string length
|
|
struct.pack_into("!BH", pkt, 1, 2 + 2 + len(topic_bytes), client.pid)
|
|
|
|
# 2. Write the packet to the socket
|
|
client.sock.write(pkt)
|
|
client._send_str(topic_bytes)
|
|
|
|
# 3. Wait for the UNSUBACK confirmation frame (0xB0) from the broker
|
|
while True:
|
|
op = client.wait_msg()
|
|
if op == 0xB0:
|
|
resp = client.sock.read(3)
|
|
assert resp[1] == pkt[2] and resp[2] == pkt[3]
|
|
return client
|
|
return client
|
|
|
|
if isinstance(topic, bytes):
|
|
topic = topic.decode('utf-8')
|
|
|
|
return client.unsubscribe(topic)
|
|
|
|
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() |