monolith/threshold

#!/usr/bin/env python
from twisted.internet import reactor
from twisted.internet.defer import Deferred
from twisted.internet.ssl import DefaultOpenSSLContextFactory
from twisted.internet.protocol import Protocol, Factory
from twisted.internet.endpoints import SSL4ClientEndpoint, TCP4ClientEndpoint, connectProtocol
from twisted.words.protocols.irc import IRCClient

from json import load, dump, loads
from sys import exit

listener = None
connections = {}
IRCPool = {}

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 IRCBot(IRCClient):
    def __init__(self, name):
        self.connected = False
        self.channels = []

        self.name = name
        instance = pool[name]

        self.nickname = instance["nickname"]
        self.realname = instance["realname"]
        self.username = instance["username"]
        self.userinfo = instance["userinfo"]
        self.fingerReply = instance["finger"]
        self.versionName = instance["version"]
        self.versionNum = None
        self.versionEnv = None
        self.sourceURL = instance["source"]
        self.autojoin = instance["autojoin"]

        self._who = {}
        self._getWho = {}
        self.wholist = {}

        self.authtype = instance["authtype"]
        if self.authtype == "ns":
            self.authpass = instance["password"]
            self.authentity = instance["authentity"]
        else:
            self.password = instance["password"]

    def refresh(self):
        instance = pool[self.name]
        if not instance["nickname"] == self.nickname:
            self.nickname = instance["nickname"]
            self.setNick(self.nickname)

        self.userinfo = instance["userinfo"]
        self.fingerReply = instance["finger"]
        self.versionName = instance["version"]
        self.versionNum = None
        self.versionEnv = None
        self.sourceURL = instance["source"]

    def privmsg(self, user, channel, msg):
        toSend = helper.isKeyword(msg)
        if toSend:
            if self.name == config["Master"][0] and channel == config["Master"][1]:
                pass
            else:
                helper.sendMaster("MATCH PRV %s (U:%s T:%s): (%s/%s) %s" % (self.name, toSend[1], toSend[2], user, channel, toSend[0]))

    def noticed(self, user, channel, msg):
        toSend = helper.isKeyword(msg)
        if toSend:
            if self.name == config["Master"][0] and channel == config["Master"][1]:
                pass
            else:
                helper.sendMaster("MATCH NOT %s (U:%s T:%s): (%s/%s) %s" % (self.name, toSend[1], toSend[2], user, channel, toSend[0]))

    def action(self, user, channel, msg):
        toSend = helper.isKeyword(msg)
        if toSend:
            if self.name == config["Master"][0] and channel == config["Master"][1]:
                pass
            else:
                helper.sendMaster("MATCH ACT %s (U:%s T:%s): (%s/%s) %s" % (self.name, toSend[1], toSend[2], user, channel, toSend[0]))

    def get(self, var):
        try:
            result = getattr(self, var)
        except AttributeError:
            result = None
        return result

    def setNick(self, nickname):
        self._attemptedNick = nickname
        self.sendLine("NICK %s" % nickname)
        self.nickname = nickname

    def alterCollidedNick(self, nickname):
        newnick = nickname + "_"
        return newnick

    def irc_ERR_NICKNAMEINUSE(self, prefix, params):
        self._attemptedNick = self.alterCollidedNick(self._attemptedNick)
        self.setNick(self._attemptedNick)

    def irc_ERR_PASSWDMISMATCH(self, prefix, params):
        log("%s: password mismatch" % self.name)
        helper.sendAll("%s: password mismatch" % self.name)

    def who(self, channel):
        channel = channel
        d = Deferred()
        if channel not in self._who:
            self._who[channel] = ([], {})
        self._who[channel][0].append(d)
        self.sendLine("WHO %s" % channel)
        return d

    def irc_RPL_WHOREPLY(self, prefix, params):
        channel = params[1]
        user = params[2]
        host = params[3]
        server = params[4]
        nick = params[5]
        status = params[6]
        realname = params[7]
        if channel not in self._who:
            return
        n = self._who[channel][1]
        n[nick] = [nick, user, host, server, status, realname]

    def irc_RPL_ENDOFWHO(self, prefix, params):
        channel = params[1]
        if channel not in self._who:
            return
        callbacks, info = self._who[channel]
        for cb in callbacks:
            cb.callback((channel, info))
        del self._who[channel]

    def got_who(self, whoinfo):
        self.wholist[whoinfo[0]] = whoinfo[1]

    def signedOn(self):
        self.connected = True
        if self.authtype == "ns":
            self.msg(self.authentity, "IDENTIFY %s" % self.nspass)
        for i in self.autojoin:
            self.join(i)

    def joined(self, channel):
        if not channel in self.channels:
            self.channels.append(channel)
        self.who(channel).addCallback(self.got_who)

    def left(self, channel):
        if channel in self.channels:
            self.channels.remove(channel)

    def kickedFrom(self, channel, kicker, message):
        if channel in self.channels:
            self.channels.remove(channel)

    def connectionLost(self, reason):
        self.connected = False
        self.channels = []
        error = reason.getErrorMessage()
        log("%s: connection lost: %s" % (self.name, error))
        helper.sendAll("%s: connection lost: %s" % (self.name, error))

    def connectionFailed(self, reason):
        self.connected = False
        self.channels = []
        error = reason.getErrorMessage()
        log("%s: connection failed: %s" % (self.name, error))
        helper.sendAll("%s: connection failed: %s" % (self.name, error))

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 saveConfig(self):
        global config
        with open("config.json", "w") as f:
            dump(config, f, indent=4)
        return

    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 sendAll(self, data):
        global connections
        for i in connections:
            connections[i].send(data)
        return

    def sendMaster(self, data):
        if config["Master"][0] in IRCPool.keys():
            IRCPool[config["Master"][0]].msg(config["Master"][1], data)

    def isKeyword(self, msg):
        message = msg.lower()
        uniqueNum = 0
        totalNum = 0
        for i in config["Keywords"]:
            if i in message:
                totalNum += message.count(i)
                message = message.replace(i, "{"+i+"}")
                uniqueNum += 1
        if totalNum == 0:
            return False
        else:
            return [message, uniqueNum, totalNum]

    def addBot(self, name):
        global IRCPool
        instance = pool[name]
        log("Started bot %s to %s:%s protocol %s nickname %s" % (name, instance["host"], instance["port"], instance["protocol"], instance["nickname"]))
        if instance["protocol"] == "plain":
            if instance["bind"] == None:
                point = TCP4ClientEndpoint(reactor, instance["host"], int(instance["port"]), timeout=int(instance["timeout"]))
                bot = IRCBot(name)
                IRCPool[name] = bot
                d = connectProtocol(point, bot)
                return
            else:
                point = TCP4ClientEndpoint(reactor, instance["host"], int(instance["port"]), timeout=int(instance["timeout"]), bindAddress=instance["bind"])
                bot = IRCBot(name)
                IRCPool[name] = bot
                d = connectProtocol(point, bot)
                return
        elif instance["protocol"] == "ssl":
            contextFactory = DefaultOpenSSLContextFactory(instance["key"].encode("utf-8", "replace"), instance["certificate"].encode("utf-8", "replace"))
            if instance["bind"] == None:
                point = SSL4ClientEndpoint(reactor, instance["host"], int(instance["port"]), contextFactory, timeout=int(instance["timeout"]))
                bot = IRCBot(name)
                IRCPool[name] = bot
                d = connectProtocol(point, bot)
                return
            else:
                point = SSL4ClientEndpoint(reactor, instance["host"], int(instance["port"]), contextFactory, timeout=int(instance["timeout"]), bindAddress=instance["bind"])
                bot = IRCBot(name)
                IRCPool[name] = bot
                d = connectProtocol(point, bot)
                return

    def addKeyword(self, keyword):
        if keyword in config["Keywords"]:
            return "EXISTS"
        else:
            for i in config["Keywords"]:
                if i in keyword or keyword in i:
                    return "ISIN"
            config["Keywords"].append(keyword)
            helper.saveConfig()
            return True

    def delKeyword(self, keyword):
        if not keyword in config["Keywords"]:
            return "NOKEY"
        config["Keywords"].remove(keyword)
        helper.saveConfig()
        return True

    def parseCommand(self, addr, authed, data):
        global pool
        data = data.strip("\n")
        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 == "rehash":
                global config
                config = helper.getConfig()
                success("Configuration rehashed successfully")

            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("Server: %s" % i)
                    for x in pool[i].keys():
                        poolMap.append("%s: %s" % (x, pool[i][x]))
                    poolMap.append("\n")
                info("\n".join(poolMap))
                return

            elif cmd == "enable":
                if length == 2:
                    if not spl[1] in pool.keys():
                        failure("Name does not exist: %s" % spl[1])
                        return
                    pool[spl[1]]["enabled"] = True
                    helper.savePool()
                    if not spl[1] in IRCPool.keys():
                        self.addBot(spl[1])
                    success("Successfully enabled bot %s" % spl[1])
                    return
                else:
                    incUsage("enable")
                    return

            elif cmd == "disable":
                if length == 2:
                    if not spl[1] in pool.keys():
                        failure("Name does not exist: %s" % spl[1])
                        return
                    pool[spl[1]]["enabled"] = False
                    helper.savePool()
                    if spl[1] in IRCPool.keys():
                        if IRCPool[spl[1]].connected == True:
                            IRCPool[spl[1]].transport.loseConnection()
                            del IRCPool[spl[1]]
                    success("Successfully disabled bot %s" % spl[1])
                    return
                else:
                    incUsage("disable")
                    return

            elif cmd == "join":
                if length == 3:
                    if not spl[1] in pool.keys():
                        failure("Name does not exist: %s" % spl[1])
                        return
                    if not spl[1] in IRCPool.keys():
                        failure("Name has no instance: %s" % spl[1])
                        return
                    IRCPool[spl[1]].join(spl[2])
                    success("Joined %s" % spl[2])
                    return
                elif length == 4:
                    if not spl[1] in pool.keys():
                        failure("Name does not exist: %s" % spl[1])
                        return
                    if not spl[1] in IRCPool.keys():
                        failure("Name has no instance: %s" % spl[1])
                        return
                    IRCPool[spl[1]].join(spl[2], spl[3])
                    success("Joined %s with key %s" % (spl[2], spl[3]))
                    return
                else:
                    incUsage("join")
                    return

            elif cmd == "part":
                if length == 3:
                    if not spl[1] in pool.keys():
                        failure("Name does not exist: %s" % spl[1])
                        return
                    if not spl[1] in IRCPool.keys():
                        failure("Name has no instance: %s" % spl[1])
                        return
                    IRCPool[spl[1]].part(spl[2])
                    success("Left %s" % spl[2])
                    return
                else:
                    incUsage("part")
                    return

            elif cmd == "get":
                if length == 3:
                    if not spl[1] in pool.keys():
                        failure("Name does not exist: %s" % spl[1])
                        return
                    if not spl[1] in IRCPool.keys():
                        failure("Name has no instance: %s" % spl[1])
                        return
                    info(str(IRCPool[spl[1]].get(spl[2])))
                    return
                else:
                    incUsage("get")
                    return

            elif cmd == "key":
                if data.startswith("key add"):
                    if not data == "key add " and not data == "key add":
                        keywordsToAdd = data[8:]
                        keywords = keywordsToAdd.split(",")
                        for keyword in keywords:
                            rtrn = self.addKeyword(keyword)
                            if rtrn == "EXISTS":
                                failure("Keyword already exists: %s" % keyword)
                            elif rtrn == "ISIN":
                                failure("Keyword already matched: %s" % keyword)
                            elif rtrn == True:
                                success("Keyword added: %s" % keyword)
                        return
                    else:
                        incUsage("key")
                        return
                elif data.startswith("key del"):
                    if not data == "key del " and not data == "key del":
                        keywordsToDel = data[8:]
                        keywords = keywordsToDel.split(",")
                        for keyword in keywords:
                            rtrn = self.delKeyword(keyword)
                            if rtrn == "NOKEY":
                                failure("Keyword does not exist: %s" % keyword)
                            elif rtrn == True:
                                success("Keyword deleted: %s" % keyword)
                        return
                    else:
                        incUsage("key")
                        return

                elif spl[1] == "show":
                    info(", ".join(config["Keywords"]))
                    return

                elif spl[1] == "master":
                    if length == 4:
                        if not spl[2] in pool.keys():
                            failure("Name does not exist: %s" % spl[2])
                            return
                        if spl[2] in IRCPool.keys():
                            if not spl[3] in IRCPool[spl[2]].channels:
                                info("Bot not on channel: %s" % spl[3])
                        config["Master"] = [spl[2], spl[3]]
                        helper.saveConfig()
                        success("Master set to %s on %s" % (spl[3], spl[2]))
                        return
                    elif length == 2:
                        info(" - ".join(config["Master"]))
                        return
                    else:
                        incUsage("key")
                        return

                else:
                    incUsage("key")
                    return

            elif cmd == "add":
                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

                    try:
                        int(spl[3])
                    except:
                        failure("Port must be an integer, not %s" % spl[3])
                        return

                    pool[spl[1]] = { "host": spl[2],
                                     "port": spl[3],
                                     "protocol": protocol,
                                     "bind": None,
                                     "timeout": 30,
                                     "nickname": spl[5],
                                     "username": config["Default"]["username"],
                                     "realname": None,
                                     "userinfo": None,
                                     "finger": None,
                                     "version": None,
                                     "source": None,
                                     "autojoin": [],
                                     "authtype": config["Default"]["authtype"],
                                     "password": config["Default"]["password"],
                                     "authentity": "NickServ",
                                     "key": config["ListenerKey"],
                                     "certificate": config["ListenerCertificate"],
                                     "enabled": config["ConnectOnCreate"],
                                   }
                    if config["ConnectOnCreate"] == True:
                        self.addBot(spl[1])
                    success("Successfully created bot")
                    self.savePool()
                    return
                else:
                    incUsage("add")
                    return

            elif cmd == "del":
                if length == 2:

                    if not spl[1] in pool.keys():
                        failure("Name does not exist: %s" % spl[1])
                        return
                    del pool[spl[1]]
                    if spl[1] in IRCPool.keys():
                        if IRCPool[spl[1]].connected == True:
                            IRCPool[spl[1]].transport.loseConnection()
                        del IRCPool[spl[1]]
                    success("Successfully removed bot")
                    self.savePool()
                    return
                else:
                    incUsage("del")
                    return

            elif cmd == "mod":
                toUnset = False
                if length == 2:
                    if not spl[1] in pool.keys():
                        failure("Name does not exist: %s" % spl[1])
                        return
                    optionMap = ["Viewing options for %s" % spl[1]]
                    for i in pool[spl[1]].keys():
                        optionMap.append("%s: %s" % (i, pool[spl[1]][i]))
                    info("\n".join(optionMap))
                    return

                elif length == 3:
                    if not spl[1] in pool.keys():
                        failure("Name does not exist: %s" % spl[1])
                        return
                    if not spl[2] in pool[spl[1]].keys():
                        failure("No such key: %s" % spl[2])
                        return
                    info("%s: %s" % (spl[2], pool[spl[1]][spl[2]]))
                    return

                elif length == 4:
                    if not spl[1] in pool.keys():
                        failure("Name does not exist: %s" % spl[1])
                        return
                    if not spl[2] in pool[spl[1]].keys():
                        failure("No such key: %s" % spl[2])
                        return
                    if spl[2] in ["port", "timeout"]:
                        try:
                            int(spl[3])
                        except:
                            failure("Value must be an integer, not %s" % spl[3])
                            return
                    if spl[2] == "protocol":
                        if not spl[3] in ["ssl", "plain"]:
                            failure("Protocol must be ssl or plain, not %s" % spl[3])
                            return
                    if spl[3] == pool[spl[1]][spl[2]]:
                        failure("Value already exists: %s" % spl[3])
                        return

                    if spl[3].lower() in ["none", "nil"]:
                        spl[3] = None
                        toUnset = True

                    if spl[2] == "authtype":
                        if not toUnset:
                            if not spl[3] in ["sp", "ns"]:
                                failure("Authtype must be sp or ns, not %s" % spl[3])
                                return
                    if spl[2] == "enabled":
                        failure("Use the enable and disable commands to manage this")
                        return
                    if spl[2] == "autojoin":
                        spl[3] = spl[3].split(",")

                    pool[spl[1]][spl[2]] = spl[3]
                    if spl[1] in IRCPool.keys():
                        IRCPool[spl[1]].refresh()
                    self.savePool()
                    if toUnset:
                        success("Successfully unset key %s on %s" % (spl[2], spl[1]))
                    else:
                        success("Successfully set key %s to %s on %s" % (spl[2], spl[3], spl[1]))
                    return

                else:
                    incUsage("mod")
                    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()
    for i in pool.keys():
        if pool[i]["enabled"] == True:
            helper.addBot(i)

    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()