truenas-stacks/mesh-mqtt-observer/app/main.py

import json
import logging
import os
import signal
import ssl
import sys
import threading
import time
from dataclasses import dataclass
from typing import Any

import orjson
import paho.mqtt.client as mqtt
import psycopg
from psycopg.rows import dict_row


STOP_EVENT = threading.Event()


def env_bool(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "y", "on"}


def env_int(name: str, default: int) -> int:
    value = os.getenv(name)
    if not value:
        return default
    return int(value)


def split_topics(value: str) -> list[str]:
    return [x.strip() for x in value.split(",") if x.strip()]


def setup_logging() -> None:
    level_name = os.getenv("LOG_LEVEL", "INFO").upper()
    level = getattr(logging, level_name, logging.INFO)

    logging.basicConfig(
        level=level,
        format="%(asctime)s [%(levelname)s] %(message)s",
        stream=sys.stdout,
    )


@dataclass
class BrokerConfig:
    source_system: str
    network_name: str | None
    modem_preset: str | None
    host: str
    port: int
    username: str | None
    password: str | None
    tls: bool
    topics: list[str]


class Database:
    def __init__(self, dsn: str):
        self.dsn = dsn

    def connect(self):
        return psycopg.connect(self.dsn, row_factory=dict_row, autocommit=True)

    def wait_until_ready(self, timeout_seconds: int = 120) -> None:
        deadline = time.time() + timeout_seconds
        last_error = None

        while time.time() < deadline:
            try:
                with self.connect() as conn:
                    with conn.cursor() as cur:
                        cur.execute("SELECT 1")
                logging.info("Postgres connection ready")
                return
            except Exception as exc:
                last_error = exc
                logging.warning("Waiting for Postgres: %s", exc)
                time.sleep(3)

        raise RuntimeError(f"Postgres not ready after {timeout_seconds}s: {last_error}")

    def init_schema(self) -> None:
        schema_path = os.path.join(os.path.dirname(__file__), "schema.sql")

        with open(schema_path, "r", encoding="utf-8") as f:
            sql = f.read()

        with self.connect() as conn:
            with conn.cursor() as cur:
                cur.execute(sql)

        logging.info("Schema initialized")

    def insert_raw(
        self,
        *,
        source_system: str,
        network_name: str | None,
        modem_preset: str | None,
        broker_host: str,
        broker_port: int,
        topic: str,
        qos: int,
        retain: bool,
        payload_raw: bytes | None,
        payload_text: str | None,
        payload_json: Any | None,
        parse_status: str,
        parser_error: str | None,
    ) -> int:
        with self.connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO raw_mqtt_events (
                        source_system,
                        network_name,
                        modem_preset,
                        broker_host,
                        broker_port,
                        topic,
                        qos,
                        retain,
                        payload_raw,
                        payload_text,
                        payload_json,
                        parse_status,
                        parser_error
                    )
                    VALUES (
                        %(source_system)s,
                        %(network_name)s,
                        %(modem_preset)s,
                        %(broker_host)s,
                        %(broker_port)s,
                        %(topic)s,
                        %(qos)s,
                        %(retain)s,
                        %(payload_raw)s,
                        %(payload_text)s,
                        %(payload_json)s,
                        %(parse_status)s,
                        %(parser_error)s
                    )
                    RETURNING id
                    """,
                    {
                        "source_system": source_system,
                        "network_name": network_name,
                        "modem_preset": modem_preset,
                        "broker_host": broker_host,
                        "broker_port": broker_port,
                        "topic": topic,
                        "qos": qos,
                        "retain": retain,
                        "payload_raw": payload_raw,
                        "payload_text": payload_text,
                        "payload_json": json.dumps(payload_json) if payload_json is not None else None,
                        "parse_status": parse_status,
                        "parser_error": parser_error,
                    },
                )
                return int(cur.fetchone()["id"])

    def insert_normalized(
        self,
        *,
        raw_event_id: int,
        source_system: str,
        network_name: str | None,
        modem_preset: str | None,
        topic: str,
        from_node_id: str | None,
        to_node_id: str | None,
        message_type: str | None,
        text_body: str | None,
        rssi: float | None,
        snr: float | None,
        payload: Any | None,
    ) -> None:
        with self.connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO normalized_events (
                        raw_event_id,
                        source_system,
                        network_name,
                        modem_preset,
                        topic,
                        from_node_id,
                        to_node_id,
                        message_type,
                        text_body,
                        rssi,
                        snr,
                        payload
                    )
                    VALUES (
                        %(raw_event_id)s,
                        %(source_system)s,
                        %(network_name)s,
                        %(modem_preset)s,
                        %(topic)s,
                        %(from_node_id)s,
                        %(to_node_id)s,
                        %(message_type)s,
                        %(text_body)s,
                        %(rssi)s,
                        %(snr)s,
                        %(payload)s
                    )
                    """,
                    {
                        "raw_event_id": raw_event_id,
                        "source_system": source_system,
                        "network_name": network_name,
                        "modem_preset": modem_preset,
                        "topic": topic,
                        "from_node_id": from_node_id,
                        "to_node_id": to_node_id,
                        "message_type": message_type,
                        "text_body": text_body,
                        "rssi": rssi,
                        "snr": snr,
                        "payload": json.dumps(payload) if payload is not None else None,
                    },
                )

    def upsert_node(self, source_system: str, node_id: str, payload: Any | None) -> None:
        with self.connect() as conn:
            with conn.cursor() as cur:
                cur.execute(
                    """
                    INSERT INTO nodes_seen (
                        source_system,
                        node_id,
                        first_seen,
                        last_seen,
                        metadata
                    )
                    VALUES (
                        %(source_system)s,
                        %(node_id)s,
                        now(),
                        now(),
                        %(metadata)s
                    )
                    ON CONFLICT (source_system, node_id)
                    DO UPDATE SET
                        last_seen = now(),
                        metadata = COALESCE(EXCLUDED.metadata, nodes_seen.metadata)
                    """,
                    {
                        "source_system": source_system,
                        "node_id": node_id,
                        "metadata": json.dumps(payload) if payload is not None else None,
                    },
                )


def try_decode_text(payload: bytes) -> str | None:
    try:
        text = payload.decode("utf-8")
    except UnicodeDecodeError:
        return None

    if not text:
        return None

    if any(ord(ch) < 9 for ch in text):
        return None

    return text


def try_decode_json(payload_text: str | None) -> tuple[Any | None, str | None]:
    if payload_text is None:
        return None, None

    stripped = payload_text.strip()

    if not stripped:
        return None, None

    if not stripped.startswith("{") and not stripped.startswith("["):
        return None, None

    try:
        return orjson.loads(stripped), None
    except Exception as exc:
        return None, str(exc)


def deep_get(data: Any, keys: list[str]) -> Any:
    if isinstance(data, dict):
        for key in keys:
            if key in data:
                return data[key]

        for value in data.values():
            result = deep_get(value, keys)
            if result is not None:
                return result

    if isinstance(data, list):
        for item in data:
            result = deep_get(item, keys)
            if result is not None:
                return result

    return None


def as_str(value: Any) -> str | None:
    if value is None:
        return None
    return str(value)


def as_float(value: Any) -> float | None:
    if value is None:
        return None

    try:
        return float(value)
    except Exception:
        return None


def infer_type(payload_json: Any, topic: str) -> str:
    explicit = deep_get(
        payload_json,
        [
            "message_type",
            "messageType",
            "type",
            "portnum",
            "portNum",
            "port_num",
            "event_type",
            "eventType",
        ],
    )

    if explicit is not None:
        text = str(explicit).lower()

        if "text" in text:
            return "text"
        if "position" in text:
            return "position"
        if "telemetry" in text:
            return "telemetry"
        if "node" in text or "user" in text:
            return "nodeinfo"

        return text

    body = deep_get(
        payload_json,
        [
            "text",
            "message",
            "body",
            "decoded_text",
            "decodedText",
            "payload_text",
            "payloadText",
            "content",
        ],
    )

    if body:
        return "text"

    lowered_topic = topic.lower()

    if "position" in lowered_topic or "/pos" in lowered_topic:
        return "position"
    if "telemetry" in lowered_topic:
        return "telemetry"
    if "nodeinfo" in lowered_topic or "node" in lowered_topic or "user" in lowered_topic:
        return "nodeinfo"

    return "unknown"


def normalize_json(payload_json: Any, topic: str) -> dict[str, Any]:
    from_node_id = as_str(
        deep_get(
            payload_json,
            [
                "from",
                "from_id",
                "fromId",
                "from_node",
                "fromNode",
                "from_node_id",
                "fromNodeId",
                "sender",
                "sender_id",
                "senderId",
                "node_id",
                "nodeId",
                "id",
            ],
        )
    )

    to_node_id = as_str(
        deep_get(
            payload_json,
            [
                "to",
                "to_id",
                "toId",
                "to_node",
                "toNode",
                "to_node_id",
                "toNodeId",
                "recipient",
                "recipient_id",
                "recipientId",
            ],
        )
    )

    text_body = as_str(
        deep_get(
            payload_json,
            [
                "text",
                "message",
                "body",
                "decoded_text",
                "decodedText",
                "payload_text",
                "payloadText",
                "content",
            ],
        )
    )

    rssi = as_float(deep_get(payload_json, ["rssi", "rx_rssi", "rxRssi"]))
    snr = as_float(deep_get(payload_json, ["snr", "rx_snr", "rxSnr"]))

    return {
        "from_node_id": from_node_id,
        "to_node_id": to_node_id,
        "message_type": infer_type(payload_json, topic),
        "text_body": text_body,
        "rssi": rssi,
        "snr": snr,
    }


class MeshMqttClient:
    def __init__(self, db: Database, config: BrokerConfig):
        self.db = db
        self.config = config

        client_id = f"mesh-mqtt-observer-{config.source_system}-{os.getpid()}"

        self.client = mqtt.Client(
            mqtt.CallbackAPIVersion.VERSION2,
            client_id=client_id,
            clean_session=True,
        )

        if config.username:
            self.client.username_pw_set(config.username, config.password or None)

        if config.tls:
            self.client.tls_set(cert_reqs=ssl.CERT_REQUIRED)
            self.client.tls_insecure_set(False)

        self.client.on_connect = self.on_connect
        self.client.on_disconnect = self.on_disconnect
        self.client.on_subscribe = self.on_subscribe
        self.client.on_message = self.on_message

    def on_connect(self, client, userdata, flags, reason_code, properties):
        logging.info(
            "[%s] Connected to MQTT broker %s:%s: %s",
            self.config.source_system,
            self.config.host,
            self.config.port,
            reason_code,
        )

        for topic in self.config.topics:
            logging.info("[%s] Subscribing to %s", self.config.source_system, topic)
            client.subscribe(topic, qos=0)

    def on_disconnect(self, client, userdata, disconnect_flags, reason_code, properties):
        logging.warning(
            "[%s] Disconnected from MQTT broker: %s",
            self.config.source_system,
            reason_code,
        )

    def on_subscribe(self, client, userdata, mid, reason_codes, properties):
        logging.info(
            "[%s] Subscribe ACK mid=%s reason_codes=%s",
            self.config.source_system,
            mid,
            reason_codes,
        )

    def on_message(self, client, userdata, msg):
        payload = bytes(msg.payload or b"")

        payload_raw = payload if env_bool("STORE_RAW_BYTES", True) else None
        payload_text = try_decode_text(payload) if env_bool("STORE_TEXT_PAYLOAD", True) else None
        payload_json, json_error = try_decode_json(payload_text)

        parse_status = "json" if payload_json is not None else "raw"

        try:
            raw_id = self.db.insert_raw(
                source_system=self.config.source_system,
                network_name=self.config.network_name,
                modem_preset=self.config.modem_preset,
                broker_host=self.config.host,
                broker_port=self.config.port,
                topic=msg.topic,
                qos=msg.qos,
                retain=msg.retain,
                payload_raw=payload_raw,
                payload_text=payload_text,
                payload_json=payload_json,
                parse_status=parse_status,
                parser_error=json_error,
            )

            if payload_json is not None:
                normalized = normalize_json(payload_json, msg.topic)

                self.db.insert_normalized(
                    raw_event_id=raw_id,
                    source_system=self.config.source_system,
                    network_name=self.config.network_name,
                    modem_preset=self.config.modem_preset,
                    topic=msg.topic,
                    from_node_id=normalized["from_node_id"],
                    to_node_id=normalized["to_node_id"],
                    message_type=normalized["message_type"],
                    text_body=normalized["text_body"],
                    rssi=normalized["rssi"],
                    snr=normalized["snr"],
                    payload=payload_json,
                )

                if normalized["from_node_id"]:
                    self.db.upsert_node(
                        self.config.source_system,
                        normalized["from_node_id"],
                        payload_json,
                    )

            logging.debug(
                "[%s] Stored topic=%s bytes=%s status=%s",
                self.config.source_system,
                msg.topic,
                len(payload),
                parse_status,
            )

        except Exception as exc:
            logging.exception(
                "[%s] Failed to store topic=%s: %s",
                self.config.source_system,
                msg.topic,
                exc,
            )

    def start(self):
        logging.info(
            "[%s] Connecting to %s:%s",
            self.config.source_system,
            self.config.host,
            self.config.port,
        )

        self.client.connect(self.config.host, self.config.port, keepalive=60)
        self.client.loop_start()

    def stop(self):
        logging.info("[%s] Stopping MQTT client", self.config.source_system)
        self.client.loop_stop()
        self.client.disconnect()


def load_configs() -> list[BrokerConfig]:
    configs: list[BrokerConfig] = []

    if env_bool("MESHTASTIC_ENABLED", True):
        configs.append(
            BrokerConfig(
                source_system="meshtastic",
                network_name=os.getenv("MESHTASTIC_NETWORK_NAME") or "Meshtastic",
                modem_preset=os.getenv("MESHTASTIC_MODEM_PRESET") or None,
                host=os.getenv("MESHTASTIC_MQTT_HOST", "mqtt.meshtastic.org"),
                port=env_int("MESHTASTIC_MQTT_PORT", 1883),
                username=os.getenv("MESHTASTIC_MQTT_USERNAME") or None,
                password=os.getenv("MESHTASTIC_MQTT_PASSWORD") or None,
                tls=env_bool("MESHTASTIC_MQTT_TLS", False),
                topics=split_topics(os.getenv("MESHTASTIC_MQTT_TOPICS", "msh/#")),
            )
        )

    if env_bool("MESHCORE_ENABLED", False):
        meshcore_host = os.getenv("MESHCORE_MQTT_HOST", "").strip()

        if meshcore_host and meshcore_host != "CHANGE_ME":
            configs.append(
                BrokerConfig(
                    source_system="meshcore",
                    network_name=os.getenv("MESHCORE_NETWORK_NAME") or "MeshCore",
                    modem_preset=os.getenv("MESHCORE_MODEM_PRESET") or None,
                    host=meshcore_host,
                    port=env_int("MESHCORE_MQTT_PORT", 1883),
                    username=os.getenv("MESHCORE_MQTT_USERNAME") or None,
                    password=os.getenv("MESHCORE_MQTT_PASSWORD") or None,
                    tls=env_bool("MESHCORE_MQTT_TLS", False),
                    topics=split_topics(os.getenv("MESHCORE_MQTT_TOPICS", "meshcore/#")),
                )
            )
        else:
            logging.warning("MeshCore is enabled but MESHCORE_MQTT_HOST is not set; skipping")

    return configs


def handle_signal(signum, frame):
    logging.info("Received signal %s; stopping", signum)
    STOP_EVENT.set()


def main() -> None:
    setup_logging()

    signal.signal(signal.SIGTERM, handle_signal)
    signal.signal(signal.SIGINT, handle_signal)

    dsn = os.getenv("DATABASE_URL")

    if not dsn:
        raise RuntimeError("DATABASE_URL is required")

    db = Database(dsn)
    db.wait_until_ready()
    db.init_schema()

    configs = load_configs()

    if not configs:
        raise RuntimeError("No MQTT brokers configured")

    clients = [MeshMqttClient(db, config) for config in configs]

    for client in clients:
        client.start()

    logging.info("mesh-mqtt-observer started with %s broker(s)", len(clients))

    try:
        while not STOP_EVENT.is_set():
            time.sleep(1)
    finally:
        for client in clients:
            client.stop()

        logging.info("mesh-mqtt-observer stopped")


if __name__ == "__main__":
    main()