screen_status_to_mqtt/mqtt_client.py

import json
import logging
import socket
import subprocess
import sys
from os import path
from time import sleep

import paho.mqtt.client as mqtt
import yaml
from yaml import SafeLoader

import app_path

logger = logging.getLogger("MqttClient")

configLocation = app_path.config_path

ha_objects = [
    {
        "component": "sensor",
        "unique_id": "display",
        "name": "Display",
        "category": "diagnostic",
        "icon": "mdi:monitor-shimmer",
        "additional_data": None,
        "init_payload": "On"
    },
    {
        "component": "binary_sensor",
        "unique_id": "monitor",
        "name": "Monitor",
        "category": "diagnostic",
        "icon": "mdi:monitor",
        "additional_data": None,
        "init_payload": "ON"
    }
]


class MqttClient:
    prefix = None
    discovery_prefix = None
    device_name = None
    device_id = None
    host = None
    port = None
    username = None
    password = None

    def __init__(self, client_name) -> None:
        super().__init__()
        self.client = mqtt.Client(client_name)
        self.client.enable_logger(logging.getLogger("MqttLib"))
        self.client.user_data_set({"self": self})

    def init_connection(self):
        if not path.exists(app_path.config_path):
            self._generate_default_conf()
        self._load_conf()
        self.client.username_pw_set(self.username, self.password)
        self.client.connect(self.host, port=self.port)
        self.client.on_connect = MqttClient.__on_connect
        self.client.loop_start()

    def disconnect(self):
        message = self.client.publish(f"{self.prefix}/mqtt", payload="offline")
        message.wait_for_publish(timeout=10)
        self.client.disconnect()

    def publish_status(self, sensor, payload):
        return self.client.publish(f"{self.prefix}/{sensor}/state", payload=payload)

    def _publish_home_assistant_discovery(self):
        for ha_object in ha_objects:
            self._send_discovery_payload_for(ha_object["component"], ha_object["unique_id"], ha_object["name"],

                                             ha_object["category"], ha_object["icon"], ha_object["additional_data"])

    def _publish_init_payload(self):
        for ha_object in ha_objects:
            self.publish_status(ha_object["unique_id"], ha_object["init_payload"])

    @staticmethod
    def __on_connect(client, userdata, flag, rc):
        logger.info("Connected to MQTT: " + mqtt.connack_string(rc))
        self: MqttClient = userdata['self']
        client.publish(f"{self.prefix}/mqtt", payload="online")
        self._publish_init_payload()
        self._publish_home_assistant_discovery()
        self._publish_init_payload()

    def _send_discovery_payload_for(self, component, unique_id, name, category, icon, additional_data=None):
        topic = f"{self.discovery_prefix}/{component}/{self.device_id}/{self.device_id}_sensor_{unique_id}/config"
        safe_name = self.device_name.lower().replace(' ', '_').encode('ascii', "ignore").decode()
        payload = {
            "state_topic": f"{self.prefix}/{unique_id}/state",
            "availability": {
                "topic": f"{self.prefix}/mqtt"
            },
            "device": {
                "identifiers": [
                    self.device_id
                ],
                "name": self.device_name
            },
            "entity_category": category,
            "icon": icon,
            "name": f"{self.device_name} {name}",
            "unique_id": f"{self.device_id} {unique_id}",
            "object_id": f"{safe_name} {unique_id}"
        }
        if additional_data:
            payload.update(additional_data)

        self.client.publish(topic, json.dumps(payload))

    def _generate_default_conf(self):
        logger.warning("Generate default conf")
        current_machine_id = subprocess.check_output('wmic csproduct get uuid').split()[1].decode("utf-8")
        default = {
            "mqtt_host": "exemple.local",
            "mqtt_port": 1883,
            "mqtt_username": "username",
            "mqtt_password": "password",
            "mqtt_prefix": socket.gethostname(),
            "discovery_prefix": "homeassistant",
            "device_name": socket.gethostname(),
            "device_id": current_machine_id
        }

        with open(app_path.config_path, "w") as stream:
            yaml.dump(default, stream)

    def _load_conf(self):
        with open(app_path.config_path, "r") as stream:
            conf = yaml.load(stream, Loader=SafeLoader)
            keys = ["mqtt_host", "mqtt_username", "mqtt_password"]

            if all(key in conf for key in keys):
                current_machine_id = subprocess.check_output('wmic csproduct get uuid').split()[1].decode("utf-8")
                self.prefix = conf.get("mqtt_prefix", "screen_status")
                self.discovery_prefix = conf.get("discovery_prefix", "homeassistant")
                self.device_name = conf.get("device_name", socket.gethostname())
                self.device_id = conf.get("device_id", current_machine_id)
                self.host = conf.get("mqtt_host")
                self.port = conf.get("mqtt_port", 1883)
                self.username = conf.get("mqtt_username")
                self.password = conf.get("mqtt_password")

            else:
                logger.critical("CONFIG FILE INVALID.")
                logger.critical("Delete the config file dans restart the service if you want to restore the default "
                                "config.")
                sys.exit(1)