🎉 Initial commit
This commit is contained in:
Jonas Bengtsson
@@ -0,0 +1,177 @@
|
||||
"""MQTT client that subscribes to an MQTTS broker and ingests telemetry.
|
||||
|
||||
Supports three broker authentication modes (see configure_tls):
|
||||
- TLS-PSK: symmetric pre-shared key, no certificates needed.
|
||||
- Client certificate (mTLS): presented when the cert files exist.
|
||||
- MQTT username/password: sent in the CONNECT packet; combinable with
|
||||
either TLS mode.
|
||||
|
||||
Persists received messages to TimescaleDB and logs locally and/or to Loki.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import ssl
|
||||
import time
|
||||
|
||||
import paho.mqtt.client as mqtt
|
||||
from paho.mqtt.enums import CallbackAPIVersion
|
||||
|
||||
from db_writer import DbWriter
|
||||
from logging_setup import setup_logging
|
||||
|
||||
logger = setup_logging()
|
||||
|
||||
BROKER_HOST = os.getenv("MQTT_BROKER_HOST")
|
||||
BROKER_PORT = int(os.getenv("MQTT_BROKER_PORT", "8883"))
|
||||
TOPIC = os.getenv("MQTT_TOPIC", "sensors/#")
|
||||
CLIENT_ID = os.getenv("MQTT_CLIENT_ID", "mqtt-ingester")
|
||||
|
||||
CA_CERT = os.getenv("MQTT_CA_CERT", "/certs/ca.crt")
|
||||
CLIENT_CERT = os.getenv("MQTT_CLIENT_CERT", "/certs/client.crt")
|
||||
CLIENT_KEY = os.getenv("MQTT_CLIENT_KEY", "/certs/client.key")
|
||||
# Empty string would be treated as a real (wrong) password by OpenSSL.
|
||||
CLIENT_KEY_PASSWORD = os.getenv("MQTT_CLIENT_KEY_PASSWORD") or None
|
||||
|
||||
PSK_IDENTITY = os.getenv("MQTT_PSK_IDENTITY", "")
|
||||
PSK_HEX = os.getenv("MQTT_PSK_HEX", "")
|
||||
|
||||
MQTT_USERNAME = os.getenv("MQTT_USERNAME", "")
|
||||
MQTT_PASSWORD = os.getenv("MQTT_PASSWORD", "")
|
||||
|
||||
writer = DbWriter()
|
||||
|
||||
|
||||
def on_connect(client, userdata, flags, reason_code, properties):
|
||||
if reason_code == 0:
|
||||
logger.info("Connected to broker %s:%s", BROKER_HOST, BROKER_PORT)
|
||||
client.subscribe(TOPIC, qos=1)
|
||||
logger.info("Subscribed to topic %s", TOPIC)
|
||||
else:
|
||||
logger.error("Connection failed with reason code %s", reason_code)
|
||||
|
||||
|
||||
def on_message(client, userdata, msg):
|
||||
topic = msg.topic
|
||||
try:
|
||||
payload = json.loads(msg.payload.decode("utf-8", errors="replace"))
|
||||
except json.JSONDecodeError:
|
||||
payload = {"raw": msg.payload.decode("utf-8", errors="replace")}
|
||||
|
||||
# Debug level: at INFO this logs (and possibly ships to Loki) once per
|
||||
# message, which dominates ingestion cost at high message rates.
|
||||
logger.debug("Received message on %s", topic, extra={"topic": topic})
|
||||
|
||||
# Enqueues only; all blocking DB work happens on the writer thread.
|
||||
writer.write(topic, payload)
|
||||
|
||||
|
||||
def on_disconnect(client, userdata, reason_code, properties):
|
||||
logger.warning("Disconnected from broker (reason %s)", reason_code)
|
||||
|
||||
|
||||
def configure_tls() -> ssl.SSLContext:
|
||||
"""Build the SSL context for the broker connection.
|
||||
|
||||
TLS-PSK takes precedence when MQTT_PSK_IDENTITY/MQTT_PSK_HEX are set;
|
||||
otherwise a certificate-based context is used. The client certificate
|
||||
is optional in the latter case, so brokers that authenticate via MQTT
|
||||
username/password (or anonymous) also work.
|
||||
"""
|
||||
if PSK_IDENTITY or PSK_HEX:
|
||||
if not (PSK_IDENTITY and PSK_HEX):
|
||||
raise RuntimeError(
|
||||
"TLS-PSK requires both MQTT_PSK_IDENTITY and MQTT_PSK_HEX to be set"
|
||||
)
|
||||
if not getattr(ssl, "HAS_PSK", False):
|
||||
raise RuntimeError(
|
||||
"TLS-PSK requested but this Python/OpenSSL build has no PSK "
|
||||
"support (ssl.HAS_PSK is False; requires Python 3.13+)"
|
||||
)
|
||||
try:
|
||||
psk = bytes.fromhex(PSK_HEX)
|
||||
except ValueError:
|
||||
raise RuntimeError("MQTT_PSK_HEX is not a valid hex string") from None
|
||||
if not psk:
|
||||
raise RuntimeError("MQTT_PSK_HEX must not be empty")
|
||||
|
||||
context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
||||
# With PSK there is no certificate to verify; the shared key
|
||||
# authenticates both sides. PSK cipher suites are TLS 1.2 only.
|
||||
context.check_hostname = False
|
||||
context.verify_mode = ssl.CERT_NONE
|
||||
context.maximum_version = ssl.TLSVersion.TLSv1_2
|
||||
context.set_ciphers("PSK")
|
||||
context.set_psk_client_callback(lambda hint: (PSK_IDENTITY, psk))
|
||||
logger.info("Using TLS-PSK authentication (identity %s)", PSK_IDENTITY)
|
||||
return context
|
||||
|
||||
context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
||||
context.load_verify_locations(cafile=CA_CERT)
|
||||
if os.path.exists(CLIENT_CERT) and os.path.exists(CLIENT_KEY):
|
||||
context.load_cert_chain(CLIENT_CERT, CLIENT_KEY, CLIENT_KEY_PASSWORD)
|
||||
logger.info("Using client certificate %s for broker auth", CLIENT_CERT)
|
||||
else:
|
||||
logger.info(
|
||||
"Client certificate not found at %s; connecting with CA verification"
|
||||
" only (broker must accept username/password or anonymous auth)",
|
||||
CLIENT_CERT,
|
||||
)
|
||||
return context
|
||||
|
||||
|
||||
def main() -> None:
|
||||
if not BROKER_HOST:
|
||||
logger.error("MQTT_BROKER_HOST is not set")
|
||||
return
|
||||
|
||||
client = mqtt.Client(
|
||||
callback_api_version=CallbackAPIVersion.VERSION2,
|
||||
client_id=CLIENT_ID,
|
||||
protocol=mqtt.MQTTv5,
|
||||
)
|
||||
|
||||
# TLS is always enabled: the stack is MQTTS-only. The mode (PSK vs
|
||||
# certificates) and the presence of a client certificate are decided
|
||||
# by configure_tls().
|
||||
client.tls_set_context(configure_tls())
|
||||
|
||||
if MQTT_USERNAME:
|
||||
# Credentials travel inside the CONNECT packet, so they work with
|
||||
# any TLS mode (PSK, cert, or CA-only) and any broker config such
|
||||
# as Mosquitto's password_file.
|
||||
client.username_pw_set(MQTT_USERNAME, MQTT_PASSWORD)
|
||||
logger.info("Using MQTT username/password authentication")
|
||||
elif not PSK_IDENTITY and not (
|
||||
os.path.exists(CLIENT_CERT) and os.path.exists(CLIENT_KEY)
|
||||
):
|
||||
logger.warning(
|
||||
"No broker authentication configured: no username/password, no PSK, "
|
||||
"and no client certificate found at %s. The broker must allow "
|
||||
"anonymous connections.",
|
||||
CLIENT_CERT,
|
||||
)
|
||||
|
||||
client.on_connect = on_connect
|
||||
client.on_message = on_message
|
||||
client.on_disconnect = on_disconnect
|
||||
|
||||
# loop_forever() reconnects automatically; this bounds the backoff so a
|
||||
# long outage does not stretch into multi-minute gaps between attempts.
|
||||
client.reconnect_delay_set(min_delay=1, max_delay=60)
|
||||
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
client.connect(BROKER_HOST, BROKER_PORT, keepalive=60)
|
||||
client.loop_forever()
|
||||
except Exception: # noqa: BLE001
|
||||
logger.exception("Connection error, retrying in 5s")
|
||||
time.sleep(5)
|
||||
finally:
|
||||
# Blocks briefly to flush any queued rows before the process exits.
|
||||
writer.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user