338 lines
11 KiB
Python
338 lines
11 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
|
|
import _thread
|
|
import gc
|
|
BACKEND_NAME = "umqtt.simple"
|
|
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 = []
|
|
self._cadata = None # Cache cert bytes to prevent heap fragmentation
|
|
|
|
# Thread safety lock for MicroPython socket reads/writes
|
|
if IS_MICROPYTHON:
|
|
self._lock = _thread.allocate_lock()
|
|
|
|
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:
|
|
gc.collect() # Clean Python heap before importing/allocating SSL
|
|
import ssl
|
|
ssl_params = self.ssl_params
|
|
|
|
if self.use_tls and ssl_params is None:
|
|
# OPTION A: If broker uses 'require_certificate false' and self-signed certs:
|
|
# Do NOT pass cadata when cert_reqs is CERT_NONE to save ~20KB of C-DRAM
|
|
ssl_params = {
|
|
"cert_reqs": ssl.CERT_NONE,
|
|
"server_hostname": self.host
|
|
}
|
|
|
|
# OPTION B: If strict CA validation IS required, load cadata ONLY with CERT_REQUIRED:
|
|
# ssl_params = {
|
|
# "cert_reqs": ssl.CERT_REQUIRED,
|
|
# "cadata": _read_file_bytes(self.cafile),
|
|
# "server_hostname": self.host
|
|
# }
|
|
|
|
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):
|
|
if IS_MICROPYTHON:
|
|
gc.collect() # Force C & Python memory cleanup right before TLS handshake
|
|
|
|
if self._client is not None:
|
|
self.close()
|
|
|
|
client = self.open()
|
|
|
|
try:
|
|
if IS_MICROPYTHON:
|
|
gc.collect() # Sweep memory right before umqtt calls ssl.wrap_socket()
|
|
with self._lock:
|
|
client.connect()
|
|
return client
|
|
|
|
client.connect(self.host, self.port, self.keepalive)
|
|
return client
|
|
except Exception as e:
|
|
print("MQTT connection failed, closing client and releasing memory.")
|
|
print("Exception:", e)
|
|
self.close()
|
|
raise
|
|
|
|
def publish(self, topic, payload, qos=2, retain=False):
|
|
client = self.open()
|
|
payload_bytes = _ensure_bytes(payload)
|
|
if IS_MICROPYTHON:
|
|
with self._lock:
|
|
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:
|
|
with self._lock:
|
|
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
|
|
import time
|
|
topic_bytes = topic if isinstance(topic, bytes) else topic.encode('utf-8')
|
|
|
|
# 1. Increment and lock the PID for THIS specific request
|
|
client.pid = (client.pid % 65535) + 1
|
|
sent_pid = client.pid # <-- Store local copy
|
|
|
|
# 2. Construct UNSUBSCRIBE packet
|
|
rem_len = 2 + 2 + len(topic_bytes)
|
|
pkt = bytearray(b"\xa2\0\0\0")
|
|
struct.pack_into("!BH", pkt, 1, rem_len, sent_pid)
|
|
|
|
# 3. Write packet to socket
|
|
client.sock.write(pkt)
|
|
client._send_str(topic_bytes)
|
|
|
|
# 4. Wait for UNSUBACK (0xB0)
|
|
start = time.time()
|
|
while time.time() - start < 3:
|
|
op = client.wait_msg()
|
|
if op == 0xB0:
|
|
resp = bytearray(3)
|
|
read_bytes = 0
|
|
while read_bytes < 3:
|
|
chunk = client.sock.read(3 - read_bytes)
|
|
if chunk:
|
|
resp[read_bytes:read_bytes + len(chunk)] = chunk
|
|
read_bytes += len(chunk)
|
|
else:
|
|
time.sleep_ms(10)
|
|
|
|
# Compare against sent_pid instead of client.pid
|
|
resp_pid = (resp[1] << 8) | resp[2]
|
|
if resp_pid != sent_pid:
|
|
print(f"[MQTT] UNSUBACK PID mismatch (expected {sent_pid}, got {resp_pid})")
|
|
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:
|
|
with self._lock:
|
|
return self._client.check_msg()
|
|
return self._client.loop(timeout=timeout)
|
|
|
|
def wait(self):
|
|
if self._client is None:
|
|
return None
|
|
if IS_MICROPYTHON:
|
|
with self._lock:
|
|
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):
|
|
"""Safely clean up socket context without causing ESP32 C panics."""
|
|
if self._client is None:
|
|
return
|
|
|
|
if IS_MICROPYTHON:
|
|
with self._lock:
|
|
try:
|
|
if hasattr(self._client, "sock") and self._client.sock:
|
|
self._client.sock.close()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
if hasattr(self._client, "sock"):
|
|
self._client.sock = None
|
|
self._client = None
|
|
gc.collect() # Immediately reclaim freed socket & mbedTLS RAM
|
|
else:
|
|
try:
|
|
self._client.disconnect()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
self._client = None
|
|
|
|
def ping(self):
|
|
"""Thread-safe PINGREQ wrapper for MicroPython."""
|
|
if self._client is None:
|
|
return
|
|
if IS_MICROPYTHON:
|
|
with self._lock:
|
|
return self._client.ping()
|
|
else:
|
|
# Paho handles keepalives automatically via loop_start/loop
|
|
pass
|
|
|
|
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() |