MqttMonitor/MqttClient.py

# -*- coding: utf-8 -*-

# Copyright (c) 2018 - 2021 Detlev Offenbach <detlev@die-offenbachs.de>
#

"""
Module implementing a PyQt wrapper around the paho MQTT client.
"""

import enum

from PyQt6.QtCore import (
    pyqtSignal, pyqtSlot, QObject, QCoreApplication, QTimer
)

import paho.mqtt.client as mqtt
from paho.mqtt.packettypes import PacketTypes
from paho.mqtt.properties import Properties

from Utilities.crypto import pwConvert


class MqttProtocols(enum.IntEnum):
    """
    Class defining the supported MQTT protocol versions.
    """
    MQTTv31 = mqtt.MQTTv31
    MQTTv311 = mqtt.MQTTv311
    MQTTv5 = mqtt.MQTTv5


class MqttClient(QObject):
    """
    Class implementing a PyQt wrapper around the paho MQTT client.
    
    @signal onConnectV3(flags, rc) emitted after the client has connected to
        the broker (MQTT v3)
    @signal onConnectV5(flags, rc, packetType, properties emitted after the
        client has connected to the broker (MQTT v5)
    @signal onDisconnectedV3(rc) emitted after the client has disconnected from
        the broker (MQTT v3)
    @signal onDisconnectedV5(rc, packetType) emitted after the client has
        disconnected from the broker (MQTT v5)
    @signal onLog(level, message) emitted to send client log data
    @signal onMessageV3(topic, payload, qos, retain) emitted after a message
        has been received by the client (MQTT v3)
    @signal onMessageV5(topic, payload, qos, retain, properties) emitted after
        a message has been received by the client (MQTT v5)
    @signal onPublish(mid) emitted after a message has been published
    @signal onSubscribeV3(mid, grantedQos) emitted after the client has
        subscribed to some topics (MQTT v3)
    @signal onSubscribeV5(mid, reasonCodes, properties) emitted after the
        client has subscribed to some topics (MQTT v5)
    @signal onUnsubscribeV3(mid) emitted after the client has unsubscribed from
        some topics (MQTT v3)
    @signal onUnsubscribeV5(mid, rc, packetType, properties) emitted after the
        client has unsubscribed from some topics (MQTT v5)
    @signal connectTimeout() emitted to indicate, that a connection attempt
        timed out
    """
    onConnectV3 = pyqtSignal(dict, int)
    onConnectV5 = pyqtSignal(dict, int, int, dict)
    onDisconnectedV3 = pyqtSignal(int)
    onDisconnectedV5 = pyqtSignal(int, int)
    onLog = pyqtSignal(int, str)
    onMessageV3 = pyqtSignal(str, bytes, int, bool)
    onMessageV5 = pyqtSignal(str, bytes, int, bool, dict)
    onPublish = pyqtSignal(int)
    onSubscribeV3 = pyqtSignal(int, tuple)
    onSubscribeV5 = pyqtSignal(int, list, dict)
    onUnsubscribeV3 = pyqtSignal(int)
    onUnsubscribeV5 = pyqtSignal(int, int, int, dict)
    
    connectTimeout = pyqtSignal()
    
    DefaultConnectTimeout = 15      # connect timeout in seconds
    
    LogDebug = 0x01
    LogInfo = 0x02
    LogNotice = 0x04
    LogWarning = 0x08
    LogError = 0x10
    LogDisabled = 0xff
    LogLevelMap = {
        mqtt.MQTT_LOG_DEBUG: LogDebug,
        mqtt.MQTT_LOG_INFO: LogInfo,
        mqtt.MQTT_LOG_NOTICE: LogNotice,
        mqtt.MQTT_LOG_WARNING: LogWarning,                        # __NO-TASK__
        mqtt.MQTT_LOG_ERR: LogError,
    }
    
    def __init__(self, clientId="", cleanSession=True, userdata=None,
                 protocol=mqtt.MQTTv311, transport="tcp", parent=None):
        """
        Constructor
        
        @param clientId ID to be used for the client
        @type str
        @param cleanSession flag indicating to start a clean session
        @type bool
        @param userdata user data
        @type any
        @param protocol version of the MQTT protocol to use
        @type int, one of mqtt.MQTTv31, mqtt.MQTTv311 or mqtt.MQTTv5
        @param transport transport to be used
        @type str, one of "tcp" or "websockets"
        @param parent reference to the parent object
        @type QObject
        """
        QObject.__init__(self, parent=parent)
        
        self.__loopStarted = False
        
        self.__connectTimeoutTimer = QTimer(self)
        self.__connectTimeoutTimer.setSingleShot(True)
        self.__connectTimeoutTimer.setInterval(
            MqttClient.DefaultConnectTimeout * 1000)
        self.__connectTimeoutTimer.timeout.connect(self.__connectTimeout)
        
        self.onConnectV3.connect(self.__connectTimeoutTimer.stop)
        self.onConnectV5.connect(self.__connectTimeoutTimer.stop)
        
        self.__cleanSession = cleanSession
        self.__protocol = protocol
        
        if protocol == MqttProtocols.MQTTv5:
            cleanSession = None
        
        self.__mqttClient = mqtt.Client(
            client_id=clientId, clean_session=cleanSession, userdata=userdata,
            protocol=int(protocol), transport=transport)
        
        self.__initCallbacks(protocol)

    def __initCallbacks(self, protocol):
        """
        Private method to initialize the MQTT callback methods.
        
        @param protocol MQTT protocol version
        @type MqttProtocols
        """
        if protocol in (MqttProtocols.MQTTv31, MqttProtocols.MQTTv311):
            self.__mqttClient.on_connect = (
                lambda client, userdata, flags, rc, properties=None:
                    self.onConnectV3.emit(flags, rc)
            )
            self.__mqttClient.on_disconnect = (
                lambda client, userdata, rc:
                    self.onDisconnectedV3.emit(rc)
            )
            self.__mqttClient.on_subscribe = (
                lambda client, userdata, mid, grantedQos, properties=None:
                    self.onSubscribeV3.emit(mid, grantedQos)
            )
            self.__mqttClient.on_unsubscribe = (
                lambda client, userdata, mid:
                    self.onUnsubscribeV3.emit(mid)
            )
            self.__mqttClient.on_message = (
                lambda client, userdata, message:
                    self.onMessageV3.emit(message.topic, message.payload,
                                          message.qos, message.retain)
            )
        else:
            self.__mqttClient.on_connect = (
                lambda client, userdata, flags, rc, properties=None:
                    self.onConnectV5.emit(
                        flags, rc.value, rc.packetType,
                        properties.json() if properties is not None else {}
                    )
            )
            self.__mqttClient.on_disconnect = self.__onDisconnectedV5
            self.__mqttClient.on_subscribe = (
                lambda client, userdata, mid, reasonCodes, properties=None:
                    self.onSubscribeV5.emit(
                        mid, reasonCodes,
                        properties.json() if properties is not None else {}
                    )
            )
            self.__mqttClient.on_unsubscribe = (
                lambda client, userdata, mid, properties, rc:
                    self.onUnsubscribeV5.emit(
                        mid, rc.value, rc.packetType,
                        properties.json() if properties is not None else {}
                    )
            )
            self.__mqttClient.on_message = (
                lambda client, userdata, message:
                    self.onMessageV5.emit(
                        message.topic, message.payload, message.qos,
                        message.retain, message.properties.json()
                    )
            )
        self.__mqttClient.on_log = (
            lambda client, userdata, level, buf:
                self.onLog.emit(level, buf)
        )
        self.__mqttClient.on_publish = (
            lambda client, userdata, mid:
                self.onPublish.emit(mid)
        )
    
    def __onDisconnectedV5(self, client, userdata, rc, properties=None):
        """
        Private method to handle the disconnect from the broker.
        
        @param client reference to the client object
        @type paho.mqtt.Client
        @param userdata user data
        @type Any
        @param rc result code or reason code
        @type int or ReasonCodes
        @param properties optional properties (defaults to None)
        @type dict (optional)
        """
        if isinstance(rc, int):
            packetType = PacketTypes.DISCONNECT
            resultCode = rc
        else:
            packetType = rc.packetType
            resultCode = rc.value
        self.onDisconnectedV5.emit(resultCode, packetType)
    
    @pyqtSlot()
    def __connectTimeout(self):
        """
        Private slot handling a failed connection attempt.
        """
        self.stopLoop()
        self.connectTimeout.emit()
    
    def setConnectionTimeout(self, timeout):
        """
        Public method to set the connection timeout value.
        
        @param timeout timeout value to be set in seconds
        @type int
        """
        self.__connectTimeoutTimer.setInterval(timeout * 1000)
    
    def setMaxInflightMessages(self, inflight=20):
        """
        Public method to set the maximum number of messages with QoS > 0 that
        can be part way through their network flow at once.
        
        @param inflight maximum number of messages in flight
        @type int
        """
        self.__mqttClient.max_inflight_messages_set(inflight)
    
    def setMaxQueuedMessages(self, queueSize=0):
        """
        Public method to set the maximum number of messages with QoS > 0 that
        can be pending in the outgoing message queue.
        
        @param queueSize maximum number of queued messages (0 = unlimited)
        @type int
        """
        self.__mqttClient.max_queued_messages_set(queueSize)
    
    def setUserCredentials(self, username, password=None):
        """
        Public method to set the user name and optionally the password.
        
        @param username user name to be set
        @type str
        @param password optional password
        @type str
        """
        self.__mqttClient.username_pw_set(username, password=password)
    
    def setUserData(self, userdata):
        """
        Public method to set the user data.
        
        @param userdata user data
        @type any
        """
        self.__mqttClient.user_data_set(userdata)
    
    # TODO: MQTTv5: add support for properties
    def setLastWill(self, topic, payload=None, qos=0, retain=False):
        """
        Public method to set the last will of the client.
        
        @param topic topic the will message should be published on
        @type str
        @param payload message to send as a will
        @type str, bytes, int or float
        @param qos quality of service level to use for the will
        @type int, one of 0, 1 or 2
        @param retain flag indicating to set as the "last known good"/retained
            message for the will topic
        @type bool
        """
        self.__mqttClient.will_set(topic, payload=payload, qos=qos,
                                   retain=retain)
    
    def clearLastWill(self):
        """
        Public method to remove a will that was previously configured with
        setLastWill().
        """
        self.__mqttClient.will_clear()
    
    def setTLS(self, caCerts=None, certFile=None, keyFile=None):
        """
        Public method to enable secure connections and set the TLS parameters.
        
        @param caCerts path to the Certificate Authority certificates file
        @type str
        @param certFile PEM encoded client certificate file
        @type str
        @param keyFile PEM encoded private key file
        @type str
        @return tuple containing a success flag and the error string of the
            paho-mqtt library
        @rtype tuple of (bool, str)
        """
        try:
            self.__mqttClient.tls_set(ca_certs=caCerts, certfile=certFile,
                                      keyfile=keyFile)
            return True, ""
        except (ValueError, FileNotFoundError) as err:
            return False, str(err)
        
        return False, "unspecific error occurred"
    
    def getProtocol(self):
        """
        Public method to get the MQTT protocol version.
        
        @return MQTT protocol version in use
        @rtype int
        """
        return self.__protocol
    
    def startLoop(self):
        """
        Public method to start the MQTT client loop.
        """
        self.__mqttClient.loop_start()
        self.__loopStarted = True
    
    def stopLoop(self):
        """
        Public method to stop the MQTT client loop.
        """
        self.__mqttClient.loop_stop()
        self.__loopStarted = False
    
    def connectToServer(self, host, port=1883, keepalive=60, bindAddress="",
                        reinit=True):
        """
        Public method to connect to a remote MQTT broker.
        
        @param host host name or IP address of the remote broker
        @type str
        @param port network port of the server host to connect to (default:
            1883, using TLS: 8883)
        @type int
        @param keepalive maximum period in seconds allowed between
            communications with the broker
        @type int
        @param bindAddress IP address of a local network interface to bind
            this client to
        @type str
        @param reinit flag indicating to reinitialize the MQTT client before
            trying to connect with the given parameters
        @type bool
        """
        # TODO: MQTTv5: add support for MQTTv5 properties
        self.__mqttClient.connect_async(
            host, port=port, keepalive=keepalive, bind_address=bindAddress,
            clean_start=self.__cleanSession)
        
        self.__connectTimeoutTimer.start()
        
        if not self.__loopStarted:
            self.startLoop()
    
    def connectToServerWithOptions(self, host, port=1883, bindAddress="",
                                   options=None):
        """
        Public method to connect to a remote MQTT broker.
        
        @param host host name or IP address of the remote broker
        @type str
        @param port network port of the server host to connect to (default:
            1883, using TLS: 8883)
        @type int
        @param bindAddress IP address of a local network interface to bind
            this client to
        @type str
        @param options dictionary containing the connection options. This
            dictionary should contain the keys "ClientId", "Keepalive",
            "CleanSession", "Username", "Password", "WillTopic", "WillMessage",
            "WillQos", "WillRetain", "TlsEnable", "TlsCaCert", "TlsClientCert",
            "TlsClientKey", "ConnectionTimeout"
        @type dict
        """
        if options:
            parametersDict = self.defaultConnectionOptions()
            parametersDict.update(options)
            
            self.setConnectionTimeout(parametersDict["ConnectionTimeout"])
            
            # step 1: set username and password
            if parametersDict["Username"]:
                if parametersDict["Password"]:
                    self.setUserCredentials(
                        parametersDict["Username"],
                        pwConvert(parametersDict["Password"], encode=False))
                else:
                    self.setUserCredentials(parametersDict["Username"])
            
            # step 2: set last will data
            if parametersDict["WillTopic"]:
                if parametersDict["WillMessage"]:
                    willMessage = parametersDict["WillMessage"]
                else:
                    # empty message to clear the will
                    willMessage = None
                self.setLastWill(parametersDict["WillTopic"],
                                 willMessage,
                                 parametersDict["WillQos"],
                                 parametersDict["WillRetain"])
            
            # step 3: set TLS parameters
            if parametersDict["TlsEnable"]:
                if (
                    parametersDict["TlsCaCert"] and
                    parametersDict["TlsClientCert"]
                ):
                    # use self signed client certificate
                    self.setTLS(caCerts=parametersDict["TlsCaCert"],
                                certFile=parametersDict["TlsClientCert"],
                                keyFile=parametersDict["TlsClientKey"])
                elif parametersDict["TlsCaCert"]:
                    # use CA certificate file
                    self.setTLS(caCerts=parametersDict["TlsCaCert"])
                else:
                    # use default TLS configuration
                    self.setTLS()
            
            # step 4: connect to server
            self.__cleanSession = parametersDict["CleanSession"]
            self.connectToServer(host, port=port,
                                 keepalive=parametersDict["Keepalive"])
        else:
            keepalive = self.defaultConnectionOptions["Keepalive"]
            self.connectToServer(host, port=port, keepalive=keepalive,
                                 bindAddress=bindAddress)
    
    @classmethod
    def defaultConnectionOptions(cls):
        """
        Class method to get a connection options dictionary with default
        values.
        
        @return dictionary containing the default connection options. It has
            the keys "ClientId", "Protocol", "ConnectionTimeout", "Keepalive",
            "CleanSession", "Username", "Password", "WillTopic", "WillMessage",
            "WillQos", "WillRetain", "TlsEnable", "TlsCaCert", "TlsClientCert",
            "TlsClientKey".
        @rtype dict
        """
        return {
            "ClientId": "ERIC7_MQTT_MONITOR_CLIENT",
            "Protocol": MqttProtocols.MQTTv311,
            "ConnectionTimeout": MqttClient.DefaultConnectTimeout,
            "Keepalive": 60,
            "CleanSession": True,
            "Username": "",
            "Password": "",
            "WillTopic": "",
            "WillMessage": "",
            "WillQos": 0,
            "WillRetain": False,
            "TlsEnable": False,
            "TlsCaCert": "",
            "TlsClientCert": "",
            "TlsClientKey": "",
        }
    
    def reconnectToServer(self):
        """
        Public method to reconnect the client with the same parameters.
        """
        self.__connectTimeoutTimer.start()
        
        self.__mqttClient.reconnect()
        
        if not self.__loopStarted:
            self.startLoop()
    
    def disconnectFromServer(self):
        """
        Public method to disconnect the client from the remote broker.
        """
        self.__connectTimeoutTimer.stop()
        
        # TODO: MQTTv5: add support for properties (?)
        # TODO: MQTTv5: add support for reason code
        self.__mqttClient.disconnect()
    
    def subscribe(self, topic, qos=0, properties=None):
        """
        Public method to subscribe to topics with quality of service.
        
        @param topic single topic to subscribe to or a tuple with a topic
            and a QoS or a list of tuples with a topic and a QoS each
        @type str or tuple of (str, int) or list of tuple of (str, int)
        @param qos quality of service
        @type int, one of 0, 1 or 2
        @param properties list of user properties to be sent with the
            subscription
        @type list of tuple of (str, str)
        @return tuple containing the result code and the message ID
        @rtype tuple of (int, int)
        """
        props = (
            self.__createPropertiesObject(PacketTypes.SUBSCRIBE, properties)
            if properties else
            None
        )
        return self.__mqttClient.subscribe(topic, qos=qos, properties=props)
    
    def unsubscribe(self, topic, properties=None):
        """
        Public method to unsubscribe topics.
        
        @param topic topic or list of topics to unsubscribe
        @type str or list of str
        @param properties list of user properties to be sent with the
            subscription
        @type list of tuple of (str, str)
        @return tuple containing the result code and the message ID
        @rtype tuple of (int, int)
        """
        props = (
            self.__createPropertiesObject(PacketTypes.SUBSCRIBE, properties)
            if properties else
            None
        )
        return self.__mqttClient.unsubscribe(topic, properties=props)
    
    # TODO: MQTTv5: add support for properties
    def publish(self, topic, payload=None, qos=0, retain=False):
        """
        Public method to publish to a topic.
        
        @param topic topic to publish to
        @type str
        @param payload data to be published
        @type str, bytes, int or float
        @param qos quality of service
        @type int, one of 0, 1 or 2
        @param retain flag indicating to set as the "last known good"/retained
            message for the topic
        @type bool
        @return message info object
        @rtype mqtt.MQTTMessageInfo
        """
        return self.__mqttClient.publish(topic, payload=payload, qos=qos,
                                         retain=retain)
    
    def __createPropertiesObject(self, packetType, properties):
        """
        Private method to assemble the MQTT v5 properties object.
        
        @param packetType type of the MQTT packet
        @type PacketTypes (= int)
        @param properties list of user properties
        @type list of tuple of (str, str)
        @return MQTT v5 properties object
        @rtype Properties
        """
        props = Properties(packetType)
        for userProperty in properties:
            props.UserProperty = tuple(userProperty)
        return props


def mqttConnackMessage(connackCode):
    """
    Module function to get the string associated with a CONNACK result.
    
    @param connackCode result code of the connection request
    @type int
    @return textual representation for the result code
    @rtype str
    """
    if connackCode == mqtt.CONNACK_ACCEPTED:
        return QCoreApplication.translate(
            "MqttConnackMessage",
            "Connection Accepted.")
    elif connackCode == mqtt.CONNACK_REFUSED_PROTOCOL_VERSION:
        return QCoreApplication.translate(
            "MqttConnackMessage",
            "Connection Refused: unacceptable protocol version.")
    elif connackCode == mqtt.CONNACK_REFUSED_IDENTIFIER_REJECTED:
        return QCoreApplication.translate(
            "MqttConnackMessage",
            "Connection Refused: identifier rejected.")
    elif connackCode == mqtt.CONNACK_REFUSED_SERVER_UNAVAILABLE:
        return QCoreApplication.translate(
            "MqttConnackMessage",
            "Connection Refused: broker unavailable.")
    elif connackCode == mqtt.CONNACK_REFUSED_BAD_USERNAME_PASSWORD:
        return QCoreApplication.translate(
            "MqttConnackMessage",
            "Connection Refused: bad user name or password.")
    elif connackCode == mqtt.CONNACK_REFUSED_NOT_AUTHORIZED:
        return QCoreApplication.translate(
            "MqttConnackMessage",
            "Connection Refused: not authorised.")
    else:
        return QCoreApplication.translate(
            "MqttConnackMessage",
            "Connection Refused: unknown reason.")


def mqttErrorMessage(mqttErrno):
    """
    Module function to get the error string associated with an MQTT error
    number.
    
    @param mqttErrno result code of a MQTT request
    @type int
    @return textual representation of the result code
    @rtype str
    """
    if mqttErrno == mqtt.MQTT_ERR_SUCCESS:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "No error.")
    elif mqttErrno == mqtt.MQTT_ERR_NOMEM:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "Out of memory.")
    elif mqttErrno == mqtt.MQTT_ERR_PROTOCOL:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "A network protocol error occurred when communicating with"
            " the broker.")
    elif mqttErrno == mqtt.MQTT_ERR_INVAL:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "Invalid function arguments provided.")
    elif mqttErrno == mqtt.MQTT_ERR_NO_CONN:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "The client is not currently connected.")
    elif mqttErrno == mqtt.MQTT_ERR_CONN_REFUSED:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "The connection was refused.")
    elif mqttErrno == mqtt.MQTT_ERR_NOT_FOUND:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "Message not found (internal error).")
    elif mqttErrno == mqtt.MQTT_ERR_CONN_LOST:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "The connection was lost.")
    elif mqttErrno == mqtt.MQTT_ERR_TLS:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "A TLS error occurred.")
    elif mqttErrno == mqtt.MQTT_ERR_PAYLOAD_SIZE:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "Payload too large.")
    elif mqttErrno == mqtt.MQTT_ERR_NOT_SUPPORTED:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "This feature is not supported.")
    elif mqttErrno == mqtt.MQTT_ERR_AUTH:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "Authorisation failed.")
    elif mqttErrno == mqtt.MQTT_ERR_ACL_DENIED:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "Access denied by ACL.")
    elif mqttErrno == mqtt.MQTT_ERR_UNKNOWN:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "Unknown error.")
    elif mqttErrno == mqtt.MQTT_ERR_ERRNO:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "Error defined by errno.")
    elif mqttErrno == mqtt.MQTT_ERR_QUEUE_SIZE:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "Message queue full.")
    else:
        return QCoreApplication.translate(
            "MqttErrorMessage",
            "Unknown error.")


def mqttLogLevelString(mqttLogLevel, isMqttLogLevel=True):
    """
    Module function to get the log level string associated with a log level.
    
    @param mqttLogLevel log level of the paho-mqtt client
    @type int
    @param isMqttLogLevel flag indicating a MQTT log level is given (if
        False it is the MqttClient variant, i.e. Debug being lowest)
    @type bool
    @return textual representation of the log level
    @rtype str
    """
    if isMqttLogLevel:
        try:
            logLevel = MqttClient.LogLevelMap[mqttLogLevel]
        except KeyError:
            return QCoreApplication.translate("MqttLogLevelString", "Unknown")
    else:
        logLevel = mqttLogLevel
    
    if logLevel == MqttClient.LogInfo:
        return QCoreApplication.translate("MqttLogLevelString", "Info")
    elif logLevel == MqttClient.LogNotice:
        return QCoreApplication.translate("MqttLogLevelString", "Notice")
    elif logLevel == MqttClient.LogWarning:
        return QCoreApplication.translate("MqttLogLevelString", "Warning")
    elif logLevel == MqttClient.LogError:
        return QCoreApplication.translate("MqttLogLevelString", "Error")
    elif logLevel == MqttClient.LogDebug:
        return QCoreApplication.translate("MqttLogLevelString", "Debug")
    elif logLevel == MqttClient.LogDisabled:
        return QCoreApplication.translate("MqttLogLevelString",
                                          "Logging Disabled")
    else:
        return QCoreApplication.translate("MqttLogLevelString", "Unknown")