monolith/threshold

#!/usr/bin/env python
from twisted.internet import reactor
from twisted.internet.ssl import DefaultOpenSSLContextFactory
from twisted.internet.protocol import Protocol, Factory
from json import load, dump, loads
from sys import exit

listener = None
connections = {}

def log(data):
    print("[LOG]", data)

def debug(data):
    print("[DEBUG]", data)

def warn(data):
    print("[WARNING]", data)

def error(data):
    print("[ERROR]", data)
    exit(1)

def sendData(addr, data):
    connections[addr].send(data)

def sendSuccess(addr, data):
    sendData(addr, "[y] " + data)

def sendFailure(addr, data):
    sendData(addr, "[n] " + data)

def sendInfo(addr, data):
    sendData(addr, "[i] " + data)

class Base(Protocol):
    def __init__(self, addr):
        self.addr = addr
        self.authed = False
        if config["UsePassword"] == False:
            self.authed = True

    def send(self, data):
        data += "\r\n"
        data = data.encode("utf-8", "replace")
        self.transport.write(data)

    def dataReceived(self, data):
        data = data.decode("utf-8", "replace")
        log("Data received from %s:%s -- %s" % (self.addr.host, self.addr.port, repr(data)))
        helper.parseCommand(self.addr, self.authed, data)

    def connectionMade(self):
        log("Connection from %s:%s" % (self.addr.host, self.addr.port))
        self.send("Hello.")

    def connectionLost(self, reason):
        global connections
        self.authed = False
        log("Connection lost from %s:%s -- %s" % (self.addr.host, self.addr.port, reason.getErrorMessage()))
        if not listener == None:
            if self.addr in connections.keys():
                del connections[self.addr]
            else:
                warn("Tried to remove a non-existant connection.")
        else:
            warn("Tried to remove a connection from a listener that wasn't running.")

class BaseFactory(Factory):
    def buildProtocol(self, addr):
        global connections
        entry = Base(addr)
        connections[addr] = entry
        return entry

    def send(self, addr, data):
        global connections
        if addr in connections.keys():
            connection = connections[addr]
            connection.send(data)
        else:
            return

class Helper(object):
    def getConfig(self):
        with open("config.json", "r") as f:
            config = load(f)
            if set(["Port", "BindAddress", "UseSSL", "UsePassword"]).issubset(set(config.keys())):
                if config["UseSSL"] == True:
                    if not set(["ListenerKey", "ListenerCertificate"]).issubset(set(config.keys())):
                        error("SSL is on but certificate or key is not defined")
                if config["UsePassword"] == True:
                    if not "Password" in config.keys():
                        error("Password authentication is on but password is not defined")
                return config
            else:
                error("Mandatory values missing from config")

    def getPool(self):
        with open("pool.json", "r") as f:
            data = f.read()
            if not data == "":
                pool = loads(data.strip())
                return pool
            else:
                return {}

    def savePool(self):
        global pool
        with open("pool.json", "w") as f:
            dump(pool, f, indent=4)
        return

    def getHelp(self):
        with open("help.json", "r") as f:
            help = load(f)
            return help

    def incorrectUsage(self, addr, mode):
        if mode == None:
            sendFailure(addr, "Incorrect usage")
            return
        if mode in help.keys():
            sendFailure(addr, "Usage: " + help[mode])
            return

    def parseCommand(self, addr, authed, data):
        global pool
        data = data.strip()
        spl = data.split()
        obj = connections[addr]

        success = lambda data: sendSuccess(addr, data)
        failure = lambda data: sendFailure(addr, data)
        info = lambda data: sendInfo(addr, data)

        incUsage = lambda mode: self.incorrectUsage(addr, mode)
        length = len(spl)
        if len(spl) > 0:
            cmd = spl[0]
        else:
            failure("No text was sent")
            return

        if authed == True:
            if cmd == "help":
                helpMap = []
                for i in help.keys():
                    helpMap.append("%s: %s" % (i, help[i]))
                info("\n".join(helpMap))
                return

            elif cmd == "pass":
                info("You are already authenticated")
                return

            elif cmd == "logout":
                obj.authed = False
                success("Logged out")
                return

            elif cmd == "list":
                poolMap = []
                for i in pool.keys():
                    poolMap.append("%s: %s" % (i, pool[i]))
                info("\n".join(poolMap))
                return

            elif cmd == "connect":
                if length == 6:

                    if spl[1] in pool.keys():
                        failure("Name already exists: %s" % spl[1])
                        return

                    protocol = spl[4].lower()

                    if not protocol in ["ssl", "plain"]:
                        incUsage("connect")
                        return

                    pool[spl[1]] = { "address": spl[2],
                                     "port": spl[3],
                                     "protocol": protocol,
                                     "nickname": spl[5],
                                   }
                    success("Successfully created bot")
                    self.savePool()
                    return
                else:
                    incUsage("connect")
                    return

            else:
                incUsage(None)
                return
        else:
            if cmd == "pass" and length == 2:
                if spl[1] == config["Password"]:
                    success("Authenticated successfully")
                    obj.authed = True
                    return
                else:
                    failure("Password incorrect")
                    obj.transport.loseConnection()
                    return
            else:
                incUsage(None)
                return

if __name__ == "__main__":
    helper = Helper()
    config = helper.getConfig()
    pool = helper.getPool()
    help = helper.getHelp()

    listener = BaseFactory()
    if config["UseSSL"] == True:
        reactor.listenSSL(config["Port"], listener, DefaultOpenSSLContextFactory(config["ListenerKey"], config["ListenerCertificate"]), interface=config["BindAddress"])
        log("Threshold running with SSL on %s:%s" % (config["BindAddress"], config["Port"]))
    else:
        reactor.listenTCP(config["Port"], listener, interface=config["BindAddress"])
        log("Threshold running on %s:%s" % (config["BindAddress"], config["Port"]))

    reactor.run()