mqtt_sniffer/mqtt_sniffer.py

import argparse
from console_bottombar import BottomBar
import getpass
import json
import mqtt
import os
import re
import time

VERSION = "0.1.0"
STARTTIME = time.time()
# TODO: Implement default values for bottombar_entries

HELPTEXT = """
F1: Get this help message
F2: Set a filter (regular expression) for the topic of a message
    Examples:
    * "/gfw/.*/videv" Get everything with "/gfw/" before "/videv" 
F9: Start / Stop logging to mqtt_sniffer.log
F12: Quit the mqtt sniffer
"""

def rx_mqtt(mc, userdata, message):
    global my_bb
    global logfile
    try:
        match = len(re.findall(my_bb.get_entry('msg_re'), message.topic)) > 0
    except re.error:
        print('No valid regular expression (%s). No filter active.' % my_bb.get_entry('msg_re'))
        match = True

    ts = time.time()-STARTTIME
    data = None
    try:
        data = json.loads(message.payload)
    except:
        pass
    try:
        data = message.payload.decode("utf-8")
    except:
        data = message.payload

    if match:
        print("%9.04f::%75s::%s" % (ts, message.topic, data))
        if my_bb.get_entry('log2file'):
            logfile.write("%9.04f::%s::%s\n" % (ts, message.topic, data))
            logfile.flush()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='This is a mqtt sniffer.')
    parser.add_argument('-f', dest='hostname', default='localhost', help='Hostname of the mqtt server')
    parser.add_argument('-p', dest='port', default=1883, help='Port of the mqtt server')
    parser.add_argument('-n', dest='no_credentials', action='store_true', help='Avoid asking for credentials')
    parser.add_argument('-u', dest='username', default=None, help='Set username for mqtt server')

    args = parser.parse_args()

    if not args.no_credentials:
        if args.username == None:
            args.username = input("Username: ")
        password = getpass.getpass(prompt='Password: ', stream=None) 
    else:
        password = None

    my_bb = BottomBar(VERSION, label='MQTT-Sniffer')
    my_bb.add_entry('help', 1, my_bb.FUNC_INFO, label='[F1] Help', infotext=HELPTEXT)
    my_bb.add_entry('msg_re', 2, my_bb.FUNC_TEXT, label='[F2] Filter')
    my_bb.add_entry('quit', 12, my_bb.FUNC_QUIT, "Quit", label='[F12]', right=True)
    my_bb.add_entry('log2file', 9, my_bb.FUNC_BOOL, label='[F9] Log2File', right=True)
    with open('mqtt_sniffer.log', 'w') as logfile:
        mc = mqtt.mqtt_client("mqtt_sniffer", args.hostname, port=args.port, username=args.username, password=password)
        mc.add_callback("#", rx_mqtt)
        my_bb.run()