healthcheck-container/hamqtt.py

import json
import logging

import paho.mqtt.client as mqtt

log = logging.getLogger(__name__)

VERSION = "2.0.0"

CONNECTIVITY = {"device_class": "connectivity"}

SENSORS = [
    ("internet_up", "binary_sensor", "Internet", CONNECTIVITY),
    ("internet_latency_ms", "sensor", "Internet latency",
     {"unit_of_measurement": "ms", "state_class": "measurement", "icon": "mdi:timer-outline"}),
    ("dns_up", "binary_sensor", "DNS", CONNECTIVITY),
    ("dns_latency_ms", "sensor", "DNS latency",
     {"unit_of_measurement": "ms", "state_class": "measurement", "icon": "mdi:dns"}),
    ("public_ip", "sensor", "Public IP", {"icon": "mdi:ip-network-outline"}),
    ("uptime_s", "sensor", "Host uptime",
     {"device_class": "duration", "unit_of_measurement": "s", "entity_category": "diagnostic"}),
]

WIFI_SENSORS = [
    ("wifi_connected", "binary_sensor", "Wi-Fi", CONNECTIVITY),
    ("wifi_ssid", "sensor", "Wi-Fi SSID", {"icon": "mdi:wifi"}),
    ("wifi_signal_dbm", "sensor", "Wi-Fi signal",
     {"device_class": "signal_strength", "unit_of_measurement": "dBm", "state_class": "measurement"}),
    ("wifi_tx_bitrate_mbps", "sensor", "Wi-Fi TX bitrate",
     {"device_class": "data_rate", "unit_of_measurement": "Mbit/s", "state_class": "measurement"}),
    ("wifi_bssid", "sensor", "Wi-Fi BSSID", {"icon": "mdi:access-point", "entity_category": "diagnostic"}),
    ("wifi_channel", "sensor", "Wi-Fi channel", {"icon": "mdi:radio-tower", "entity_category": "diagnostic"}),
    ("wifi_ap_count", "sensor", "Wi-Fi networks in range",
     {"icon": "mdi:access-point-network", "state_class": "measurement"}),
]


def flatten(snapshot):
    link = snapshot.get("wifi_link") or {}
    scan = snapshot.get("wifi_scan") or {}
    return {
        "internet_up": snapshot.get("internet_up"),
        "internet_latency_ms": snapshot.get("internet_latency_ms"),
        "dns_up": snapshot.get("dns_up"),
        "dns_latency_ms": snapshot.get("dns_latency_ms"),
        "public_ip": snapshot.get("public_ip"),
        "uptime_s": snapshot.get("uptime_s"),
        "wifi_connected": link.get("connected"),
        "wifi_ssid": link.get("ssid"),
        "wifi_bssid": link.get("bssid"),
        "wifi_signal_dbm": link.get("signal_dbm"),
        "wifi_tx_bitrate_mbps": link.get("tx_bitrate_mbps"),
        "wifi_channel": link.get("channel"),
        "wifi_ap_count": scan.get("count"),
    }


class HomeAssistantPublisher:
    def __init__(self, cfg):
        self.cfg = cfg
        self.node = f"healthcheck_{cfg.device_id}"
        base = f"healthcheck/{cfg.device_id}"
        self.state_topic = f"{base}/state"
        self.availability_topic = f"{base}/availability"
        self.status_topic = f"{cfg.mqtt_discovery_prefix}/status"
        self._last_payload = None

        self.client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2, client_id=self.node)
        if cfg.mqtt_username:
            self.client.username_pw_set(cfg.mqtt_username, cfg.mqtt_password)
        self.client.will_set(self.availability_topic, "offline", retain=True)
        self.client.on_connect = self._on_connect
        self.client.on_message = self._on_message
        self.client.on_disconnect = self._on_disconnect

    def start(self):
        self.client.reconnect_delay_set(min_delay=2, max_delay=60)
        self.client.connect_async(self.cfg.mqtt_host, self.cfg.mqtt_port, keepalive=60)
        self.client.loop_start()
        log.info("MQTT: connecting to %s:%s as %s", self.cfg.mqtt_host, self.cfg.mqtt_port, self.node)

    def publish_state(self, snapshot):
        payload = flatten(snapshot)
        if payload == self._last_payload:
            return
        self._last_payload = payload
        self.client.publish(self.state_topic, json.dumps(payload), retain=True)

    def _on_connect(self, client, userdata, flags, reason_code, properties):
        if reason_code.is_failure:
            log.warning("MQTT: connection refused: %s", reason_code)
            return
        log.info("MQTT: connected")
        client.subscribe(self.status_topic)
        self._announce()
        if self._last_payload is not None:
            client.publish(self.state_topic, json.dumps(self._last_payload), retain=True)

    def _on_disconnect(self, client, userdata, flags, reason_code, properties):
        log.warning("MQTT: disconnected: %s", reason_code)

    def _on_message(self, client, userdata, message):
        if message.topic == self.status_topic and message.payload.decode(errors="ignore") == "online":
            log.info("MQTT: Home Assistant came online, re-announcing")
            self._announce()

    def _announce(self):
        device = {
            "identifiers": [self.node],
            "name": self.cfg.device_name,
            "model": "healthcheck-container",
            "manufacturer": "healthcheck",
            "sw_version": VERSION,
        }
        sensors = SENSORS + (WIFI_SENSORS if self.cfg.wifi_interface else [])
        for key, component, name, extra in sensors:
            if component == "binary_sensor":
                template = (
                    f"{{{{ 'None' if value_json.{key} is none "
                    f"else ('ON' if value_json.{key} else 'OFF') }}}}"
                )
            else:
                template = f"{{{{ value_json.{key} }}}}"
            payload = {
                "name": name,
                "unique_id": f"{self.node}_{key}",
                "state_topic": self.state_topic,
                "value_template": template,
                "availability_topic": self.availability_topic,
                "device": device,
                **extra,
            }
            topic = f"{self.cfg.mqtt_discovery_prefix}/{component}/{self.node}/{key}/config"
            self.client.publish(topic, json.dumps(payload), retain=True)
        self.client.publish(self.availability_topic, "online", retain=True)