Init
This commit is contained in:
+218
@@ -0,0 +1,218 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user