isengard-bot/main.py

#!/usr/bin/env python3

from systemd.journal import JournalHandler
from getpass import getpass
from argparse import ArgumentParser
import random
import threading
import socket
import time
import signal
import logging

from bot import MUCBot
from data import DataStore
from localServer import LocalServer

def cleanExit():
    global localservthread
    localservthread.pleaseStop = True
    time.sleep(1)
    # stop tcp server
    sockfd = socket.socket()
    port = 12346
    sockfd.connect(('127.0.0.1', port))
    sockfd.send(b'')
    sockfd.close()
    time.sleep(1)

    exit(0)


def signalHandler(sig, frame):
    log.info("Exiting...")
    cleanExit()

if __name__ == '__main__':

    # Register SIGINT signal handler
    signal.signal(signal.SIGINT, signalHandler)

    # Setup the command line arguments.
    parser = ArgumentParser()

    # Output verbosity options.
    parser.add_argument("-q", "--quiet", help="set logging to ERROR",
                        action="store_const", dest="loglevel",
                        const=logging.ERROR, default=logging.INFO)
    parser.add_argument("-d", "--debug", help="set logging to DEBUG",
                        action="store_const", dest="loglevel",
                        const=logging.DEBUG, default=logging.INFO)

    # JID and password options.
    parser.add_argument("-j", "--jid", dest="jid",
                        help="JID to use")
    parser.add_argument("-p", "--password", dest="password",
                        help="password to use")
    parser.add_argument("-r", "--rooms", dest="rooms",
                        help="MUC rooms to join separated by comma")
    parser.add_argument("-n", "--nick", dest="nick",
                        help="MUC nickname")

    args = parser.parse_args()

    # Setup logging.
    logging.basicConfig(level=args.loglevel,
                        format='%(levelname)-8s %(message)s')
    log = logging.getLogger(__name__)
    log.addHandler(JournalHandler())


    # Setup the MUCBot and register plugins. Note that while plugins may
    # have interdependencies, the order in which you register them does
    # not matter.

    jid = args.jid
    password = args.password
    rooms = args.rooms
    nick = args.nick

    if not jid:
        jid = "isengard@a-lec.org"
    if not password:
        password = getpass()
    if not rooms:
        rooms = "tests-isengard@salons.a-lec.org"
    if not nick:
        nick = "Isengard"

    ERRORS = 0
    datastore = None
    localservthread = None

    while True:

        xmpp = MUCBot(jid, password, rooms.split(","), nick)
        xmpp.register_plugin('xep_0030') # Service Discovery
        xmpp.register_plugin('xep_0045') # Multi-User Chat
        xmpp.register_plugin('xep_0199') # XMPP Ping

        # Create buffer
        if not datastore:
            datastore = DataStore(xmpp)
        xmpp.datastore = datastore

        # Launch local server
        if not localservthread:
            localservthread = LocalServer(datastore)
            localservthread.start()
            time.sleep(1)

        # Connect to the XMPP server and start processing XMPP stanzas.
        xmpp.connect()

        # Check local server status
        if not localservthread.is_alive():
            exit(1)

        try:
            log.info("Launching XMPP connection until disconnection (ERRORS: %s)" % ERRORS)
            xmpp.loop.run_until_complete(xmpp.disconnected)
        except KeyboardInterrupt:
            cleanExit()
        except Exception as e:
            log.error("ERROR IN MAIN : " + str(e))
            if ERRORS >= 3:
                log.critical("3 consecutive errors. Aborting...")
                cleanExit()
            ERRORS += 1
            del xmpp