monolith/core/bot.py

from twisted.internet.protocol import ReconnectingClientFactory
from twisted.words.protocols.irc import IRCClient
from twisted.internet.defer import Deferred
from twisted.internet.task import LoopingCall, deferLater
from twisted.internet import reactor

from string import digits
from random import randint

from modules import userinfo
from modules import counters as count
from modules import monitor

from core.relay import sendRelayNotification

import main
from utils.logging.log import *
from utils.logging.send import *

from twisted.internet.ssl import DefaultOpenSSLContextFactory

def deliverRelayCommands(relay, relayCommands, user=None, stage2=None):
    # where relay is a dictionary extracted from the Network object
    keyFN = main.certPath+main.config["Key"]
    certFN = main.certPath+main.config["Certificate"]
    contextFactory = DefaultOpenSSLContextFactory(keyFN.encode("utf-8", "replace"),
            certFN.encode("utf-8", "replace"))
    bot = IRCBotFactory(None, relay, relayCommands, user, stage2)
    rct = reactor.connectSSL(relay["host"],
            int(relay["port"]),
            bot, contextFactory)

class IRCRelay(IRCClient):
    def __init__(self, relay, relayCommands, user, stage2):
        self.connected = False
        self.buffer = ""
        if user == None:
            self.user = relay["user"]
        else:
            self.user = user
        password = relay["password"]
        self.nickname = self.user
        self.realname = self.user
        self.username = self.user
        self.password = self.user+":"+password

        self.relayCommands = relayCommands
        self.relay = relay
        self.stage2 = stage2

    def parsen(self, user):
        step = user.split("!")
        nick = step[0]
        if len(step) == 2:
            step2 = step[1].split("@")
            ident, host = step2
        else:
            ident = nick
            host = nick

        return [nick, ident, host]

    def privmsg(self, user, channel, msg):
        nick, ident, host = self.parsen(user)
        if nick[0] == main.config["Tweaks"]["ZNC"]["Prefix"]:
            nick = nick[1:]
            if nick in self.relayCommands.keys():
                sendAll("[%s] %s -> %s" % (self.relay, nick, msg))

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

    def signedOn(self):
        self.connected = True
        log("signed on as a relay: %s" % self.relay)
        if main.config["Notifications"]["Connection"]:
            sendRelayNotification("Relay", {"type": "conn", "status": "connected"}) 
        for i in self.relayCommands.keys():
            for x in self.relayCommands[i]:
                self.msg(main.config["Tweaks"]["ZNC"]["Prefix"]+i, x)
        if not self.stage2 == None: # [["user", {"sasl": ["message1", "message2"]}], []]
            if not len(self.stage2) == 0:
                user = self.stage2[0].pop(0)
                commands = self.stage2[0].pop(0)
                del self.stage2[0]
                deliverRelayCommands(self.relay, commands, user, self.stage2)
        deferLater(reactor, 1, self.transport.loseConnection)
        return

class IRCBot(IRCClient):
    def __init__(self, name, relay):
        self.connected = False
        self.channels = []
        self.net = "".join([x for x in name if not x in digits])
        if self.net == "":
            error("Network with all numbers: %s" % name)
        self.buffer = ""
        self.name = name
        alias = relay["alias"]

        self.nickname = alias["nick"]
        self.realname = alias["realname"]
        self.username = alias["nick"]+"/"+relay["net"]
        self.password = relay["password"]
        self.userinfo = None
        self.fingerReply = None
        self.versionName = None
        self.versionNum = None
        self.versionEnv = None
        self.sourceURL = None

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

        self._names = {}

    def parsen(self, user):
        step = user.split("!")
        nick = step[0]
        if len(step) == 2:
            step2 = step[1].split("@")
            if len(step2) == 2:
                ident, host = step2
            else:
                ident = nick
                host = nick
        else:
            ident = nick
            host = nick

        return [nick, ident, host]

    def privmsg(self, user, channel, msg):
        nick, ident, host = self.parsen(user)
        # ZNC adds messages when no users are in the channel, messing up tracking
        #userinfo.editUser(self.net, channel, nick, user)
        count.event(self.net, "msg")
        #event = None
        #if self.nickname.lower() in msg.lower():
        #    event = "highlight"
        monitor.event(self.net, self.name, {"type": "msg", "nick": nick, "ident": ident, "host": host, "target": channel, "message": msg})

    def noticed(self, user, channel, msg):
        nick, ident, host = self.parsen(user)
        #userinfo.editUser(self.net, channel, nick, user)
        count.event(self.net, "notice")

        monitor.event(self.net, self.name, {"type": "notice", "nick": nick, "ident": ident, "host": host, "target": channel, "message": msg})

    def action(self, user, channel, msg):
        nick, ident, host = self.parsen(user)
        #userinfo.editUser(self.net, channel, nick, user)
        count.event(self.net, "action")

        monitor.event(self.net, self.name, {"type": "action", "nick": nick, "ident": ident, "host": host, "target": channel, "message": msg})

    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 nickChanged(self, olduser, newnick):
        oldnick, ident, host = self.parsen(olduser)
        userinfo.renameUser(self.net, oldnick, olduser, newnick, newnick+"!"+ident+"@"+host)
        self.nickname = newnick
        count.event(self.net, "selfnick")

    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)
        sendAll("%s: password mismatch" % self.name)

    def who(self, 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.append([nick, user, host, server, status, realname])
        count.event(self.net, "whoreply")
        monitor.event(self.net, self.name, {"type": "who", "nick": nick, "ident": user, "host": host, "realname": realname, "target": channel, "server": server, "status": status})

    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):
        userinfo.initialUsers(self.net, whoinfo[0], whoinfo[1])

    def sanit(self, data):
        if len(data) >= 1:
            if data[0] in ["!", "~", "&", "@", "%", "+"]:
                data = data[1:]
                return data
            return data
        else:
            return False

    def names(self, channel):
        d = Deferred()
        if channel not in self._names:
            self._names[channel] = ([], [])
        self._names[channel][0].append(d)
        self.sendLine("NAMES %s" % channel)
        return d

    def irc_RPL_NAMREPLY(self, prefix, params):
        channel = params[2]
        nicklist = params[3].split(' ')
        if channel not in self._names:
            return
        n = self._names[channel][1]
        n.append(nicklist)

    def irc_RPL_ENDOFNAMES(self, prefix, params):
        channel = params[1]
        if channel not in self._names:
            return
        callbacks, namelist = self._names[channel]
        for cb in callbacks:
            cb.callback((channel, namelist))
        del self._names[channel]

    def got_names(self, nicklist):
        newNicklist = []
        for i in nicklist[1]:
            for x in i:
                f = self.sanit(x)
                if f:
                    newNicklist.append(f)
        userinfo.initialNames(self.net, nicklist[0], newNicklist)

    #twisted sucks so i have to do this to actually get the user info
    def irc_JOIN(self, prefix, params):
        """
        Called when a user joins a channel.
        """
        nick = prefix.split('!')[0]
        channel = params[-1]
        if nick == self.nickname:
            self.joined(channel)
        else:
            self.userJoined(prefix, channel)

    def irc_PART(self, prefix, params):
        """
        Called when a user leaves a channel.
        """
        nick = prefix.split('!')[0]
        channel = params[0]
        if len(params) >= 2:
            message = params[1]
        else:
            message = None
        if nick == self.nickname:
            self.left(prefix, channel, message)
        else:
            self.userLeft(prefix, channel, message)

    def irc_QUIT(self, prefix, params):
        """
        Called when a user has quit.
        """
        nick = prefix.split('!')[0]
        if nick == self.nickname:
            self.botQuit(prefix, params[0])
        else:
            self.userQuit(prefix, params[0])

    def irc_NICK(self, prefix, params):
        """
        Called when a user changes their nickname.
        """
        nick = prefix.split('!', 1)[0]
        if nick == self.nickname:
            self.nickChanged(prefix, params[0])
        else:
            self.userRenamed(prefix, params[0])

    def irc_KICK(self, prefix, params):
        """
        Called when a user is kicked from a channel.
        """
        #kicker = prefix.split('!')[0]
        channel = params[0]
        kicked = params[1]
        message = params[-1]
        if kicked.lower() == self.nickname.lower():
            # Yikes!
            self.kickedFrom(channel, prefix, message)
        else:
            self.userKicked(kicked, channel, prefix, message)

    def irc_TOPIC(self, prefix, params):
        """
        Someone in the channel set the topic.
        """
        #user = prefix.split('!')[0]
        channel = params[0]
        newtopic = params[1]
        self.topicUpdated(prefix, channel, newtopic)
    #END hacks

    def signedOn(self):
        self.connected = True
        log("signed on: %s" % self.name)
        if main.config["Notifications"]["Connection"]:
            sendRelayNotification(self.name, {"type": "conn", "status": "connected"}) 
        count.event(self.net, "signedon")

    def joined(self, channel):
        if not channel in self.channels:
            self.channels.append(channel)
        self.names(channel).addCallback(self.got_names)
        self.who(channel).addCallback(self.got_who)
        count.event(self.net, "selfjoin")
        if self.name == main.config["Master"][0] and channel == main.config["Master"][1]:
            for i in range(len(main.masterbuf)):
                self.msg(channel, main.masterbuf.pop(0))
                main.saveConf("masterbuf")
        lc = LoopingCall(self.who, channel)
        self._getWho[channel] = lc
        intrange = main.config["Tweaks"]["Delays"]["WhoRange"]
        minint = main.config["Tweaks"]["Delays"]["WhoLoop"]
        interval = randint(minint, minint+intrange)
        lc.start(interval)

    def botLeft(self, channel):
        if channel in self.channels:
            self.channels.remove(channel)
        if channel in self._getWho.keys():
            lc = self._getWho[channel]
            lc.stop()
            del self._getWho[channel]
        userinfo.delChannel(self.net, channel)

    def left(self, user, channel, message):
        nick, ident, host = self.parsen(user)
        count.event(self.net, "selfpart")
        monitor.event(self.net, self.name, {"type": "part", "target": channel, "message": message})
        monitor.event(self.net, self.name, {"type": "part", "self": True, "target": channel, "message": message})
        self.botLeft(channel)

    def kickedFrom(self, channel, kicker, message):
        nick, ident, host = self.parsen(kicker)
        if channel in self.channels:
            self.channels.remove(channel)
        count.event(self.net, "selfkick")
        monitor.event(self.net, self.name, {"type": "kick", "nick": nick, "ident": ident, "host": host, "target": channel, "message": message})
        monitor.event(self.net, self.name, {"type": "kick", "self": True, "nick": nick, "ident": ident, "host": host, "target": channel, "message": message})
        self.botLeft(channel)

    def userJoined(self, user, channel):
        nick, ident, host = self.parsen(user)
        userinfo.addUser(self.net, channel, nick, user)
        count.event(self.net, "join")
        monitor.event(self.net, self.name, {"type": "join", "nick": nick, "ident": ident, "host": host, "target": channel})

    def userLeft(self, user, channel, message):
        nick, ident, host = self.parsen(user)
        userinfo.delUser(self.net, channel, nick, user)
        count.event(self.net, "part")
        monitor.event(self.net, self.name, {"type": "part", "nick": nick, "ident": ident, "host": host, "target": channel, "message": message})

    def userQuit(self, user, quitMessage):
        nick, ident, host = self.parsen(user)
        userinfo.delUserByNetwork(self.net, nick, user)
        count.event(self.net, "quit")

        monitor.event(self.net, self.name, {"type": "quit", "nick": nick, "ident": ident, "host": host, "message": quitMessage})

    def botQuit(self, user, quitMessage):
        nick, ident, host = self.parsen(user)
        userinfo.delUserByNetwork(self.net, nick, user)
        count.event(self.net, "selfquit")

        monitor.event(self.net, self.name, {"type": "quit", "nick": nick, "ident": ident, "host": host, "message": quitMessage})
        monitor.event(self.net, self.name, {"type": "quit", "self": True, "nick": nick, "ident": ident, "host": host, "message": quitMessage})

    def userKicked(self, kickee, channel, kicker, message):
        nick, ident, host = self.parsen(kicker)
        userinfo.editUser(self.net, channel, nick, kicker)
        userinfo.delUserByNick(self.net, channel, kickee)
        count.event(self.net, "kick")

        monitor.event(self.net, self.name, {"type": "kick", "nick": nick, "ident": ident, "host": host, "target": channel, "message": message, "user": kickee})

    def userRenamed(self, oldname, newname):
        nick, ident, host = self.parsen(oldname)
        userinfo.renameUser(self.net, nick, oldname, newname, newname+"!"+ident+"@"+host)
        count.event(self.net, "nick")
        monitor.event(self.net, self.name, {"type": "nick", "nick": nick, "ident": ident, "host": host, "user": newname})

    def topicUpdated(self, user, channel, newTopic):
        nick, ident, host = self.parsen(user)
        userinfo.editUser(self.net, channel, nick, user)
        count.event(self.net, "topic")

        monitor.event(self.net, self.name, {"type": "topic", "nick": nick, "ident": ident, "host": host, "target": channel, "message": newTopic})

    def modeChanged(self, user, channel, toset, modes, args):
        nick, ident, host = self.parsen(user)
        userinfo.editUser(self.net, channel, nick, user)
        count.event(self.net, "mode")
        argList = list(args)
        modeList = [i for i in modes]
        for a, m in zip(argList, modeList):
            monitor.event(self.net, self.name, {"type": "mode", "nick": nick, "ident": ident, "host": host, "target": channel, "modes": m, "modeargs": a})

class IRCBotFactory(ReconnectingClientFactory):
    def __init__(self, name, relay=None, relayCommands=None, user=None, stage2=None):
        if not name == None:
            self.name = name
            self.net = "".join([x for x in self.name if not x in digits])
        else:
            self.name = "Relay to "+relay["net"]+relay["id"]
        self.client = None
        self.maxDelay = main.config["Tweaks"]["Delays"]["MaxDelay"]
        self.initialDelay = main.config["Tweaks"]["Delays"]["InitialDelay"]
        self.factor = main.config["Tweaks"]["Delays"]["Factor"]
        self.jitter = main.config["Tweaks"]["Delays"]["Jitter"]

        self.relay, self.relayCommands, self.user, self.stage2 = relay, relayCommands, user, stage2

    def buildProtocol(self, addr):
        if self.relay == None:
            entry = IRCBot(self.name, self.relay)
            main.IRCPool[self.name] = entry
        else:
            entry = IRCRelay(self.relay, self.relayCommands, self.user, self.stage2)

        self.client = entry
        return entry

    def clientConnectionLost(self, connector, reason):
        if not self.relay:
            userinfo.delNetwork(self.net, self.client.channels)
        if not self.client == None:
            self.client.connected = False
            self.client.channels = []
        error = reason.getErrorMessage()
        log("%s: connection lost: %s" % (self.name, error))
        if not self.relay:
            sendAll("%s: connection lost: %s" % (self.name, error))
        if main.config["Notifications"]["Connection"]:
            sendRelayNotification(self.name, {"type": "conn", "status": "lost", "reason": error})
        if not self.relay:
            self.retry(connector)
        #ReconnectingClientFactory.clientConnectionLost(self, connector, reason)

    def clientConnectionFailed(self, connector, reason):
        if not self.client == None:
            self.client.connected = False
            self.client.channels = []
        error = reason.getErrorMessage()
        log("%s: connection failed: %s" % (self.name, error))
        if not self.relay:
            sendAll("%s: connection failed: %s" % (self.name, error))
        if main.config["Notifications"]["Connection"]:
            sendRelayNotification(self.name, {"type": "conn", "status": "failed", "reason": error})
        if not self.relay:
            self.retry(connector)
        #ReconnectingClientFactory.clientConnectionFailed(self, connector, reason)