mqtt-ingester/mqtt-client/main.py

"""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()