mqttmon

#!/usr/bin/env python3

# This file is part of mqttmon.
#
# mqttmon is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2015, 2024 Dominik Kriegner <dominik.kriegner@gmail.com>

"""MQTT monitor application for monitoring MQTT broker messages."""

from __future__ import annotations

import argparse
import collections
import configparser
import curses
import curses.panel
import itertools
import sys
import time
from dataclasses import dataclass
from typing import Any, Optional

from paho.mqtt import client as mqtt_client


class TopicList(dict):
    """Dictionary for organizing MQTT topics and messages.

    A class extending a Python dictionary to hold the information about different
    topics received from the broker. For every topic a associated MessageList is
    created. The topics are sorted in a Tree structure for later processing.

    Parameters
    ----------
    name : str
        Name of the topic list
    path : str, optional
        Parent path for this topic list
    """
    def __init__(self, name: str, path: Optional[str] = None):
        """Initialize the topic list."""
        super().__init__()
        self.name = '/'.join((path, name)) if path else name
        self.msglist = MessageList(name)

    def addSubTopic(self, topicName: str, path: str) -> None:
        """Add a new topic to the TopicList.

        Parameters
        ----------
        topicName : str
            Name of the topic to be added
        path : str
            Path for the new topic
        """
        super().__setitem__(topicName, TopicList(topicName, path))

    def addMessage(self, topic: str, msg: Any) -> None:
        """Add a new message to a topic.

        Parameters
        ----------
        topic : str
            Topic to add the message to
        msg : Any
            Message object to be added
        """
        self.msglist.append(msg)
        if topic:
            lt = topic.split('/', 1)
            if lt[0] not in super().keys():
                self.addSubTopic(lt[0], self.name)
            if len(lt) > 1:
                super().__getitem__(lt[0]).addMessage(lt[1], msg)
            else:
                super().__getitem__(lt[0]).addMessage(None, msg)

    def getMessages(self, mode: str) -> list:
        """Get message list based on specified mode.

        Parameters
        ----------
        mode : str
            Type of message list to return

        Returns
        -------
        list
            Message list in the requested format
        """
        if mode == 'topics':
            return self.listTopics()
        elif mode == 'inplace':
            lm = self.lastMessage()
            if lm is None:
                return []
            else:
                return self.lastMessage()
        elif mode == 'continuous' or 'topic':
            return self.msglist.lines
        else:
            raise Exception('unknown Message list mode')

    def lastMessage(self) -> Optional[list]:
        """Return list of last messages per topic.

        Returns
        -------
        list or None
            List of tuples (indentation, line text, curses attribute) for display,
            or None if no messages
        """
        if len(self) == 0:
            if len(self.msglist) > 0:
                return self.msglist[-1].listLines()
            else:
                return None
        else:
            ret = []
            for t in super().keys():
                lines = super().__getitem__(t).lastMessage()
                if lines:
                    ret += lines
            return ret

    def listTopics(self) -> list:
        """Return formatted list of topics.

        Returns
        -------
        list
            List of tuples (indentation, line text, curses attribute) for display
        """
        ret = []
        ret.append((self, 0, self.name, curses.A_BOLD))
        mlist = self.msglist
        if len(self) > 0:
            desc = '{} subtopic{}, '.format(len(self),
                                            's' if len(self) > 0 else '')
        else:
            desc = ''
        desc += '{} msgs'.format(len(mlist))
        if len(mlist) > 0:
            desc += ', last update {}'.format(mlist[-1].timestr)
        ret.append((self, 1, desc, curses.A_NORMAL))
        for topic in sorted(self):
            ret += super().__getitem__(topic).listTopics()
        return ret


class MessageList(collections.deque):
    """List for storing and displaying MQTT messages.

    Extends collections.deque to provide text line formatting capabilities
    for displaying messages on the screen.

    Parameters
    ----------
    name : str
        Name of the message list
    maxlen : int, optional
        Maximum length for message history, by default 512
    """
    def __init__(self, name: str, maxlen: int = 512):
        """Initialize the message list."""
        self.name = name
        self.lines = collections.deque(maxlen=maxlen*2)
        super().__init__(maxlen=maxlen)

    def append(self, msg: Any) -> None:
        """Append a message to the list.

        Parameters
        ----------
        msg : Any
            Message object to append
        """
        self.lines += msg.listLines()
        super().append(msg)


class Message:
    """Object for storing and formatting MQTT messages.

    Holds message content, topic, timestamp and formatting methods
    for display.

    Parameters
    ----------
    msg : bytes
        Raw message content
    topic : str
        Topic message was received on
    retain : bool
        Message retain flag
    qos : int
        QoS level for message
    """
    _timefmt = "%H:%M:%S"
    _longtimefmt = _timefmt + "%d/%m %Y"

    def __init__(self, msg: bytes, topic: str, retain: bool, qos: int):
        """Initialize message with content and metadata."""
        self.timestamp = time.time()
        self.timestr = time.strftime(self._timefmt,
                                     time.localtime(self.timestamp))
        self.topic = topic
        self.message = msg
        self.retain = retain
        self.qos = qos

    def listLines(self) -> list:
        """Return formatted message lines for display.

        Returns
        -------
        list
            List of tuples (indentation, line text, curses attribute)
        """
        ret = [(self, 0, self.topic, curses.A_BOLD), ]
        try:
            for line in self.message.decode('utf-8').splitlines():
                ret.append((self, 1, line, curses.A_NORMAL))
        except UnicodeDecodeError:
            ret.append((self, 1, 'msg not decodable to UTF-8',
                        curses.A_UNDERLINE))
        return ret

    def listDetails(self) -> list:
        """Return detailed message information.

        Returns
        -------
        list
            List of tuples (indentation, line text, curses attribute)
        """
        ret = self.listLines()
        ret.append((self, 0, '----', curses.A_BOLD))
        tstamp = time.strftime(self._timefmt, time.localtime(self.timestamp))
        ret.append((self, 1, 'message received at {}'.format(tstamp),
                    curses.A_NORMAL))
        ret.append((self, 1, 'retain flag: {}; QOS: {}'.format(self.retain,
                                                               self.qos),
                    curses.A_NORMAL))
        return ret


class Communicator:
    """Client for communicating with MQTT broker.

    Parameters
    ----------
    address : str
        MQTT broker address
    port : int
        TCP port for connection
    user : str, optional
        Username for authentication
    passwd : str, optional
        Password for authentication
    userdata : dict, optional
        Data to pass to callbacks
    protocol : int, optional
        MQTT protocol version to use (3 for v3.1, 4 for v3.1.1, 5 for v5.0)
    """
    def __init__(self, address: str, port: int, user: Optional[str] = None,
                passwd: Optional[str] = None, userdata: Optional[dict] = None,
                protocol: int = 4):
        """Initialize connection to MQTT broker."""
        self.topics: list[str] = []
        version = mqtt_client.CallbackAPIVersion.VERSION2
        self.client = mqtt_client.Client(version, 'mqttstatus', protocol=protocol)
        if user:
            self.client.username_pw_set(user, password=passwd)
        self.client.on_connect = self.on_connect
        self.client.on_message = self.on_message
        self.client.on_subscribe = self.on_subscribe
        self.client.user_data_set(userdata)
        self.client.connect(address, port, 60)

    def on_connect(self, client: Any, userdata: dict[str, Any],
                  flags: dict[str, Any], reason_code: mqtt_client.ReasonCodes, properties = None) -> None:
        """Handle successful connection to broker.

        Parameters
        ----------
        client : Any
            Client instance for callback
        userdata : dict
            User data passed to callback
        flags : dict
            Response flags from broker
        reason_code : ReasonCodes
            Connection result code
        properties : Any, optional
            Connection properties
        """
        if 'StatusWindow' in userdata:
            cursesTUI.setStatusMessage(userdata['StatusWindow'],
                                       'Connected with result code '+str(reason_code))

    def subscribe(self, *args: str) -> None:
        """Subscribe to one or more topics.

        Parameters
        ----------
        *args : str
            Topics to subscribe to
        """
        if len(args) == 0:
            self.client.subscribe("#")
            self.topics.append('#')
        else:
            for t in args:
                self.client.subscribe(t)
                self.topics.append(t)

    def on_subscribe(self, client: Any, userdata: dict[str, Any],
                    mid: int, reason_code: list[mqtt_client.ReasonCodes], properties = None) -> None:
        """Handle subscribe confirmation from broker.

        Parameters
        ----------
        client : Any
            Client instance for callback
        userdata : dict
            User data passed to callback
        mid : int
            Message ID
        reason_code : list[ReasonCodes]
            QoS values for subscriptions
        properties : Any, optional
            Subscription properties
        """
        if 'StatusWindow' in userdata:
            t = ' '.join(self.topics)
            cursesTUI.setStatusMessage(
                userdata['StatusWindow'],
                'successfully subscribed topic: {}'.format(t))

    def on_message(self, client: Any, userdata: dict[str, Any], message: Any) -> None:
        """Handle received messages from broker.

        Parameters
        ----------
        client : Any
            Client instance for callback
        userdata : dict
            User data passed to callback
        message : Any
            MQTT message with topic, payload, qos, retain, properties
        """
        # add message to the corresponding message list
        m = Message(message.payload, message.topic, message.retain, message.qos)
        userdata['TopicList'].addMessage(message.topic, m)
        # increase message counter
        userdata['nmsg'] += 1
        # show Message on screen
        if 'StatusWindow' in userdata:
            cursesTUI.setStatusMessage(
                userdata['StatusWindow'],
                'nmsg: {:4d};'.format(userdata['nmsg']), start=1, end=12)
            cursesTUI.setStatusMessage(
                userdata['StatusWindow'],
                'New msg in {topic} at {time}'.format(topic=m.topic,
                                                      time=m.timestr))


class cursesTUI:
    """Text user interface using curses.

    Main class for MQTT monitor with window management, broker connection
    handling and display updates.

    Parameters
    ----------
    stdscr : Any
        Standard screen from curses.wrapper
    args : argparse.Namespace
        Command line arguments
    """
    statwidth = 2  # height of the status window area (minimum 2)

    def __init__(self, stdscr: Any, args: argparse.Namespace):
        """Initialize windows and broker connection."""
        self.scroll = 0  # integer constant to save the scroll offset
        self.ispaused = False  # update mode paused flag
        self.updateMode = 'inplace'  # can be 'continuous': as recv msg
                                     #        'inplace': latest msg/topic
                                     #        'topics': overview of topics
                                     #        'topic': msges from single topic
                                     #        'msg': details about message
        self.lines = []  # line buffer for shown message list
        self.cursory = 0  # cursor position to select messages

        self.mainWindow = stdscr
        self.initMain()
        self.topicList = TopicList(args.topic)
        self.userdict = {'TopicList': self.topicList,
                         'MainWindow': self.mainWindow,
                         'MessageWindow': self.msgWindow,
                         'StatusWindow': self.statusWindow,
                         'nmsg': 0}

        # start communication
        self.com = Communicator(args.brokeraddr, args.port, args.username,
                                args.passwd, self.userdict, args.protocol)
        self.com.subscribe(args.topic)
        self.mainLoop(args.timeout)

    def initMain(self):
        """Initialize main curses windows."""
        self.width_y, self.width_x = self.mainWindow.getmaxyx()

        # message window
        self.msgWindow = curses.newwin(self.width_y - self.statwidth,
                                       self.width_x)
        self.clearMsgWindow()
        wy, wx = self.msgWindow.getmaxyx()
        self.textheight = wy - 2
        self.contentwidth = wx - 2

        # status window
        self.statusWindow = curses.newwin(self.statwidth, self.width_x,
                                          self.width_y - self.statwidth, 0)
        self.redrawStatus()

    def clearMsgWindow(self, title: str = 'Messages') -> None:
        """Clear and redraw message window border.

        Parameters
        ----------
        title : str, optional
            Window title to display
        """
        self.msgWindow.erase()
        self.msgWindow.border(0, 0, 0, 0)
        self.msgWindow.addstr(0, 2, ' {} '.format(title))
        self.msgWindow.move(1, 1)
        self.scroll = 0
        self.ispaused = False
        self.lines = []

    def redrawStatus(self) -> None:
        """Redraw status window with keyboard shortcuts."""
        h, w = self.statusWindow.getmaxyx()
        if w < 58:
            raise Exception("Window is too narrow, needs at least 58 width!")
        self.statusWindow.addstr(1, 1, 'q-quit;')
        self.statusWindow.addstr(1, 9, 'msg: ')
        if self.updateMode == 'continuous':
            self.statusWindow.addstr(1, 14, 'i-inplace')
            self.statusWindow.addstr(1, 24, 'c-continuous', curses.A_BOLD)
            self.statusWindow.addstr(1, 37, 't-topics;')
        elif self.updateMode == 'inplace':
            self.statusWindow.addstr(1, 14, 'i-inplace', curses.A_BOLD)
            self.statusWindow.addstr(1, 24, 'c-continuous')
            self.statusWindow.addstr(1, 37, 't-topics;')
        else:
            self.statusWindow.addstr(1, 14, 'i-inplace')
            self.statusWindow.addstr(1, 24, 'c-continuous')
            self.statusWindow.addstr(1, 37, 't-topics', curses.A_BOLD)
            self.statusWindow.addstr(1, 44, ';')
        if self.ispaused:
            self.statusWindow.addstr(1, 47, 'p-(un)pause', curses.A_BOLD)
        else:
            self.statusWindow.addstr(1, 47, 'p-(un)pause')

    def resize(self) -> None:
        """Handle terminal window resize event."""
        old_width_y, old_width_x = (self.width_y, self.width_x)
        self.width_y, self.width_x = self.mainWindow.getmaxyx()
        self.msgWindow.resize(self.width_y - self.statwidth, self.width_x)
        self.resizeMsgWindow(old_width_y, old_width_x)
        self.statusWindow.resize(self.statwidth, self.width_x)
        self.statusWindow.mvwin(self.width_y - self.statwidth, 0)
        self.update()

    def resizeMsgWindow(self, old_width_y: int, old_width_x: int) -> None:
        """Handle message window resize.

        Parameters
        ----------
        old_width_y : int
            Previous window height
        old_width_x : int
            Previous window width
        """
        old_textheight, old_contentwidth = (self.textheight, self.contentwidth)
        wy, wx = self.msgWindow.getmaxyx()
        self.textheight = wy - 2
        self.contentwidth = wx - 2
        # vertical change
        step = 1 if (self.width_y - old_width_y) > 0 else -1
        for cl in range(old_width_y+step, self.width_y+step, step):
            if cl > old_width_y:
                self.msgWindow.move(old_textheight+1, 1)
                self.msgWindow.insertln()
            elif self.textheight > len(self.lines):
                self.msgWindow.move(self.textheight+1, 1)
                self.msgWindow.deleteln()
            else:
                self.msgWindow.move(1, 1)
                self.msgWindow.deleteln()
        # horizontal change
        step = 1 if (self.width_x - old_width_x) > 0 else -1
        for cc in range(old_width_x+step, self.width_x+step, step):
            if cc > old_width_x:
                for py in range(1, 1+self.textheight):
                    self.msgWindow.insch(py, old_contentwidth+1, ' ')
            else:
                for py in range(1, 1+self.textheight):
                    self.msgWindow.delch(py, self.contentwidth+1)
        # fix border
        self.msgWindow.border(0, 0, 0, 0)

    def mainLoop(self, timeout: int) -> None:
        """Run main program loop.

        Parameters
        ----------
        timeout : int
            Input check timeout in ms
        """
        self.msgWindow.timeout(timeout)
        curses.curs_set(2)
        while True:
            # call the communication loop of the MQTT client
            try:
                self.com.client.loop(timeout=timeout/1000/10)
            except UnicodeDecodeError:
                self.setStatusMessage(self.statusWindow,
                                      'UnicodeDecodeError in MQTT client')
            # check key-events
            c = self.msgWindow.getch()
            if c == ord('q'):
                sys.exit()
            elif c == ord('c'):
                self.clearMsgWindow()
                self.updateMode = 'continuous'
            elif c == ord('i'):
                self.clearMsgWindow()
                self.updateMode = 'inplace'
            elif c == ord('t'):
                self.clearMsgWindow(title='Topic list')
                self.updateMode = 'topics'
            elif c == ord('p'):
                self.ispaused = not self.ispaused
            elif c in (curses.KEY_UP, 65):
                self.cursory -= 1
            elif c in (curses.KEY_DOWN, 66):
                self.cursory += 1
            elif c in (curses.KEY_ENTER, 10):
                if isinstance(self.shownobjects[self.cursory], Message):
                    self.clearMsgWindow(title='Message details')
                    self.updateMode = 'msg'
                elif isinstance(self.shownobjects[self.cursory], TopicList):
                    self.clearMsgWindow(title='Topic messages')
                    self.updateMode = 'topic'
            elif c == curses.KEY_RESIZE:
                self.resize()
            self.redrawStatus()
            self.updateMsgList()
            self.update()

    def updateMsgList(self) -> None:
        """Update displayed message list based on mode."""
        self.msgWindow.move(1, 1)
        window = self.msgWindow
        # create message list
        if not self.ispaused:
            if self.updateMode in ('topic', 'msg') and self.lines == []:
                o = self.shownobjects[self.cursory]
                if isinstance(o, Message):
                    self.lines = o.listDetails()
                elif isinstance(o, TopicList):
                    self.lines = o.getMessages(self.updateMode)
            elif self.updateMode in ('inplace', 'topics', 'continuous'):
                self.lines = self.topicList.getMessages(self.updateMode)
        # handle scrolling
        if len(self.lines) < self.textheight:
            self.scroll = 0
        nlines = min(self.textheight, len(self.lines))
        if self.cursory >= nlines:
            self.scroll -= self.cursory - nlines + 1
            if self.scroll < 0:
                self.scroll = 0
            self.cursory = nlines - 1 if nlines > 0 else 0
        elif self.cursory < 0:
            self.scroll += -self.cursory
            self.cursory = 0
        if self.scroll > abs(len(self.lines) - self.textheight):
            self.scroll = len(self.lines) - self.textheight
        lstart = max(0, len(self.lines) - self.textheight - self.scroll)
        scrollinfo = ' {}-{}/{} '.format(
            lstart,
            lstart + min(len(self.lines), self.textheight),
            len(self.lines))
        self.msgWindow.addstr(self.textheight+1,
                              self.contentwidth - len(scrollinfo),
                              scrollinfo)
        # output message list
        self.msgWindow.move(1, 1)
        shownlines = itertools.islice(self.lines, lstart, lstart+nlines)
        self.shownobjects = []
        for c, (o, i, line_text, a) in enumerate(shownlines):
            self.shownobjects.append(o)
            self.addLineStr(window, line_text, self.contentwidth-i, 1+i, attr=a)
        # show cursor position
        self.msgWindow.move(self.cursory+1, 1)

    def update(self) -> None:
        """Refresh screen content."""
        self.statusWindow.refresh()
        self.msgWindow.refresh()

    @staticmethod
    def setStatusMessage(window: Any, string: str, start: int = 14, end: int = -1,
                         attr: int = curses.A_NORMAL) -> None:
        """Set status message in status bar.

        Parameters
        ----------
        window : Any
            Target curses window
        string : str
            Message to display
        start : int, optional
            Starting position, by default 14
        end : int, optional
            Ending position (-1 for line end), by default -1
        attr : int, optional
            Curses text attributes, by default normal
        """
        wy, wx = window.getmaxyx()
        window.move(0, start)
        if end == -1:
            end = wx - 2
            window.clrtoeol()
        for line in string.splitlines():
            window.addnstr(line, end-start, attr)

    @staticmethod
    def addMultiLineStr(window: Any, string: str, start: int = 1, end: int = -1,
                       attr: int = curses.A_NORMAL) -> None:
        """Add multi-line string to curses window.

        Parameters
        ----------
        window : Any
            Target curses window
        string : str
            Text to display
        start : int, optional
            Indent for each line, by default 1
        end : int, optional
            Line end position (-1 for window end), by default -1
        attr : int, optional
            Curses text attributes, by default normal
        """
        y, x = window.getyx()
        wy, wx = window.getmaxyx()
        for line in string.splitlines():
            cursesTUI.addLineStr(window, line, wy-start+end, start, attr=attr)

    @staticmethod
    def addLineStr(window: Any, line: str, nchar: int, start: int = 1,
                  attr: int = curses.A_NORMAL) -> None:
        """Add single line string to curses window.

        Parameters
        ----------
        window : Any
            Target curses window
        line : str
            Text to display
        nchar : int
            Number of characters to write
        start : int, optional
            Starting position, by default 1
        attr : int, optional
            Curses text attributes, by default normal
        """
        y, x = window.getyx()
        window.move(y, 1)
        window.addstr(' '*(start-1), curses.A_NORMAL)
        window.move(y, start)
        window.addnstr(line, nchar, attr)
        # clear the rest of the line without destroying the window border
        y, x = window.getyx()
        # this complicated treatment here is necessary for special character
        # functionality
        window.addstr(' '*(nchar-(x-start)), curses.A_NORMAL)
        y, x = window.getyx()
        window.move(y+1, start)


if __name__ == '__main__':
    conf_parser = argparse.ArgumentParser(add_help=False)
    conf_parser.add_argument('-c', '--config', metavar='conffile', type=str,
                             default='', help='specify config file')
    args, remaining_argv = conf_parser.parse_known_args()
    defaults = {
        'port': 1883,
        'passwd': None,
        'username': None,
        'topic': '#',
        'timeout': 100,
        'brokeraddr': None,
        'protocol': 4
        }
    if args.config:
        config = configparser.ConfigParser()
        config.read([args.config])
        for item, val in config.items("Defaults"):
            defaults[item] = val

    parser = argparse.ArgumentParser(parents=[conf_parser],
                                     description='curses TUI for MQTT message'
                                     ' monitoring')
    parser.set_defaults(**defaults)
    parser.add_argument('-p', '--port', metavar='port', type=int,
                        help='TCP port for the MQTT broker')
    parser.add_argument('-u', '--user', dest='username', metavar='username',
                        type=str, help='user name for the broker connection'
                        ' (empty if not needed)')
    parser.add_argument('-P', '--pass', metavar='passwd', dest='passwd',
                        type=str, help='password for the broker connection'
                        ' (empty if not needed)')
    parser.add_argument('-t', '--topic', metavar='topic', type=str,
                        help='topics to subscribe')
    parser.add_argument('-v', '--protocol', type=int, choices=[3,4,5],
                       help='MQTT protocol version (3=v3.1, 4=v3.1.1, 5=v5.0)')
    parser.add_argument('brokeraddr', metavar='brokeraddress', type=str,
                        help='Address of the MQTT broker', nargs='?')
    args = parser.parse_args(remaining_argv)
    if not args.brokeraddr:
        print('Broker-address or config file must be specified!')
        sys.exit()

    curses.wrapper(cursesTUI, args)