simple-mqtt-exporter.py

#!/usr/bin/env python3
import argparse
import importlib
import json
import os
import re
import sys
import time

import paho.mqtt.client as mqtt
import prometheus_client as prom


def smart_float(value):
    if isinstance(value, bool):
        if value:
            return 1.0
        else:
            return 0.0
    if isinstance(value, str):
        if value in (
            "on",
            "ON",
            "On",
            "true",
            "TRUE",
            "True",
            "yes",
            "YES",
            "Yes",
            "Online",
            "online",
        ):
            return 1.0
        if value in (
            "off",
            "OFF",
            "Off",
            "false",
            "FALSE",
            "False",
            "no",
            "NO",
            "No",
            "Offline",
            "offline",
        ):
            return 0.0
    return float(value)


def get_field(content, field):
    result = content
    for f in field.split("."):
        if isinstance(result, dict):
            result = result.get(f, {})
    return smart_float(result)


def regex_match(string):
    for topic in list(config.mqtt_topics.keys()):
        p = re.compile("^" + topic + "$")
        if p.search(string):
            # add regex based match to config.mqtt_topics as plain topic
            config.mqtt_topics[string] = config.mqtt_topics[topic]
            topic_init(string, config.mqtt_topics[topic])
            return True
    return False


def on_connect(client, userdata, flags, rc, properties=None):
    codes = [
        "Connection successful",
        "Connection refused – incorrect protocol version",
        "Connection refused – invalid client identifier",
        "Connection refused – server unavailable",
        "Connection refused – bad username or password",
        "Connection refused – not authorised",
    ]
    if rc != 0:
        if hasattr(mqtt, "CallbackAPIVersion"):
            print(rc)
        elif rc > 0 and rc < 6:
            print(codes[rc])
        else:
            print(f"Bad connection, unknown return code: {rc}")
        os._exit(1)


def on_message(client, userdata, msg):
    global data_received, success, error
    if config.debug:
        print(f"{msg.topic}: {msg.payload}")
    if msg.topic in config.mqtt_topics or regex_match(msg.topic):
        try:
            payload = str(msg.payload.decode("utf-8", "strict"))
            success[msg.topic] += 1
            received_messages.labels(status="success", topic=msg.topic).set(
                success[msg.topic]
            )
        except Exception as e:
            print(f"{type(e).__name__}: {str(e)} while decoding topic {msg.topic}")
            error[msg.topic] += 1
            received_messages.labels(status="error", topic=msg.topic).set(
                error[msg.topic]
            )
            return
        data_received = True
        # When the config is a list, we need to use the JSON fields in the message
        if isinstance(config.mqtt_topics[msg.topic], list):
            try:
                content = json.loads(payload)
            except Exception as e:
                print(
                    f"{type(e).__name__}: {str(e)} while decoding json topic {msg.topic}"
                )
                error[msg.topic] += 1
                received_messages.labels(status="error", topic=msg.topic).set(
                    error[msg.topic]
                )
                return
            for item in config.mqtt_topics[msg.topic]:
                field = item["field"]
                try:
                    value = get_field(content, field)
                except Exception as e:
                    print(
                        f"{type(e).__name__}: {str(e)} while decoding topic {msg.topic} field {field}"
                    )
                    error[msg.topic] += 1
                    received_messages.labels(status="error", topic=msg.topic).set(
                        error[msg.topic]
                    )
                    continue
                gauges[msg.topic + ":" + field].set(value)
        else:
            try:
                value = smart_float(payload)
            except Exception as e:
                print(
                    f"{type(e).__name__}: {str(e)} while decoding topic {msg.topic} field {field}"
                )
                error[msg.topic] += 1
                received_messages.labels(status="error", topic=msg.topic).set(
                    error[msg.topic]
                )
                return
            gauges[msg.topic].set(value)
        if not msg.retain:
            updated.set(time.time())


def topic_init(t, v):
    global success, error, gauges, parents
    success[t] = 0
    error[t] = 0
    parts = t.split("/")
    if isinstance(v, list):
        items = v
        sep = ":"
    else:
        items = [v]
        sep = ""
    for i in items:
        field = i.get("field", "")
        topic = t + sep + field
        name = i.get("name")
        if not name:
            name = parts[-1:][0]
        description = i.get("help")
        if not description:
            description = t + sep + field
        labels = i.get("labels", {})
        labels["topic"] = t
        if field:
            labels["field"] = field
        if len(parts) > 1 and not labels.get("sensor"):
            labels["sensor"] = parts[1]
        if len(parts) > 2 and not labels.get("device"):
            labels["device"] = parts[2]
        if not name in parents:
            parents[name] = prom.Gauge(name, description, labels.keys())
        try:
            gauges[topic] = parents[name].labels(**labels)
        except ValueError as e:
            print(
                f"{type(e).__name__} while adding gauge for topic {topic}: name = {name}, labels = {labels}, error = {str(e)}"
            )
            sys.exit(1)


def mqtt_init():
    if hasattr(mqtt, "CallbackAPIVersion"):
        client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
    else:
        client = mqtt.Client()
    if hasattr(config, "mqtt_username") and hasattr(config, "mqtt_password"):
        client.username_pw_set(config.mqtt_username, config.mqtt_password)
    client.on_connect = on_connect
    client.on_message = on_message
    try:
        client.connect(config.mqtt_broker)
    except (TimeoutError, OSError) as e:
        print(f"{type(e).__name__} connecting to {config.mqtt_broker}: {e}")
        sys.exit(1)
    if hasattr(config, "mqtt_twc_topic") and not hasattr(config, "mqtt_topic"):
        client.subscribe(config.mqtt_twc_topic)
    else:
        client.subscribe(config.mqtt_topic)
    client.loop_start()
    return client


if __name__ == "__main__":
    sys.stdout.reconfigure(line_buffering=True)
    sys.stderr.reconfigure(line_buffering=True)
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", "--config", help="config file to load", default="config")
    args = parser.parse_args()
    config = importlib.import_module(args.config)
    success = {}
    error = {}
    parents = {}
    gauges = {}
    for t, v in config.mqtt_topics.items():
        # Do not initialize regex topics
        if not "*" in t:
            topic_init(t, v)
    up = prom.Gauge("up", "client status")
    updated = prom.Gauge("updated", "data last updated in epoch")
    received_messages = prom.Gauge(
        "received_messages",
        "received messages per topic and status",
        ["status", "topic"],
    )
    client = mqtt_init()
    prom.start_http_server(config.http_port)

    fresh = time.time()
    stale = 0
    startup = True
    while True:
        data_received = False
        if startup:
            time.sleep(10)
        else:
            time.sleep(getattr(config, "sleep", 10))
        if data_received:
            up.set(1)
            fresh = time.time()
            stale = 0
            startup = False
        elif stale < 3:
            stale += 1
        else:
            stale += 1
            up.set(0)
            if time.time() - fresh > 600:
                print(f"Exiting, mqtt state = {client._state} for too long")
                sys.exit(1)