irc.py

import socket, os, threading, re, asynchat, traceback, ssl
import plugins

class User(str):
    def __init__(self, nick):
        super(User, self).__init__(self, nick)
        self.ident = ""
        self.host = ""

class PluginManager(object):
    def __init__(self, server):
        self.server = server
        self.plugins = {}
        self.prnt = server.prnt
        self.handlers = { 
            "join" : [ ],
            "part" : [ ],
            "quit" : [ ],
            "message" : [ ],
            "connect" : [ ],
            "disconnect": [ ],
            "action" : [ ],
            "nick" : [ ],
            "kick" : [ ],
            "mode" : [ ],
            "command" : { },
        }
    
    def handle(self, event, handler_func, handler_commands=None):
        event = event.lower()
        if event == "command":
            if handler_commands is not None:
                for command in handler_commands:
                    self.handlers[event][command.lower()] = handler_func
        else:
            if self.handlers.get(event, None) is None:
                self.handlers[event] = []
                
            self.handlers[event].append(handler_func)
    
    def unhandle(self, event, handler_func, handler_commands=None):
        event = event.lower()
        if event == "command":
            handlers = {}
            for command, handler in self.handlers[event]:
                if handler != handler_func and not command.lower() in handler_commands:
                    handlers.insert(0, handler)
                    
            self.handlers[event] = handlers
        else:
            handlers = []
            for handler in self.handlers[event]:
                if handler != handler_func:
                    handlers.insert(0, handler)
                    
            self.handlers[event] = handlers
        
    def event(self, eventName, *args):
        eventName = eventName.lower()
        if self.handlers.get(eventName, None) is None: return
        
        if eventName == "command":
            handler = self.handlers[eventName].get(args[2].lower(), None)
            if handler is None:
                return #no such command
                
            handler(*args)
        else:
            for handler in self.handlers[eventName]:
                handler(*args)
    
    def loadedPlugin(self, pluginName):
        return pluginName.lower() in self.plugins
    
    def pluginExists(self, pluginName):
        return plugins.getPlugin(pluginName.lower()) is not None
        
    def loadPlugin(self, pluginName):
        plugins.refresh()
        if not self.loadedPlugin(pluginName.lower()):
            if not self.pluginExists(pluginName.lower()):
                return
            
            pClass = plugins.getPlugin(pluginName.lower())
            
            self.plugins[pluginName.lower()] = pClass(self.server)
    
    def getPlugin(self, pluginName): 
        try:
            return self.plugins[pluginName.lower()]
        except KeyError:
            return None
    
    def reloadPlugin(self, pluginName):
        errs = {}
        
        err = self.unloadPlugin(pluginName.lower())
        if err is not None: 
            errs["unload"] = err
            
        err = self.loadPlugin(pluginName.lower())
        if err is not None: 
            errs["load"] = err
        
        return errs
    
    def unloadPlugin(self, pluginName):
        if self.loadedPlugin(pluginName.lower()):
            inst = self.plugins[pluginName.lower()]
            for event, eList in self.handlers.items():
                if event == "command":
                    newList = {}
                    for cmd,func in eList.items():
                        if func.__self__ == inst:
                            destructor = getattr(inst, "destroy", None)
                            if destructor is not None:
                                destructor()
                        else:
                            newList[cmd] = func
                    self.handlers[event] = newList
                else:
                    newList = []
                    for func in eList:
                        if func.__self__ == inst:
                            destructor = getattr(inst, "destroy", None)
                            if destructor is not None:
                                destructor()
                        else:
                            newList.append(func)
                    self.handlers[event] = newList
                            
            del self.plugins[pluginName.lower()]
        else:
            err = "Plugin "+pluginName+" is not loaded."
            return err
    
    def loadAllPlugins(self):
        errs = {}
        plugins = [k for k in self.server.config["plugins"].split(",") if k != '']
        for plugin in plugins:
            err = self.loadPlugin(plugin.lower())
            if err is not None:
                errs[plugin] = err
                
        return errs
    
    def unloadAllPlugins(self):
        errs = {}
        for plugin in self.plugins.copy().keys():
            err = self.unloadPlugin(plugin.lower())
            if err is not None:
                errs[plugin] = err
            
        return errs

class IRC(asynchat.async_chat):
    def __init__(self, config):
        asynchat.async_chat.__init__(self)
        self.pluginManager = PluginManager(self)
        self.config = config
        self.handle = self.pluginManager.handle
        self.unhandle = self.pluginManager.unhandle
        self.handlers = self.pluginManager.handlers
        self.plugins = self.pluginManager.plugins
        self.getPlugin = self.pluginManager.getPlugin
        errs = self.pluginManager.loadAllPlugins()
                
        self.userlist = { }
        self.data = ""
        self.set_terminator("\r\n")
           
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if self.config["use_ssl"].lower().startswith("y"):
            try:
                self.ssl_sock = ssl.wrap_socket(self.sock)
                self.ssl_sock.connect((self.config["host"], int(self.config["port"])))
                self.set_socket(self.ssl_sock)
            except ssl.SSLError, error:
                self.prnt('SSL Error connecting to server. (Are you using the right port?)')
                raise ssl.SSLError(error)
            except socket.error, error:
                print('There was an error connecting to %s' % address)
                return
        else:
            try:
                self.sock.connect((self.config["host"], int(self.config["port"])))
            except socket.error, error:
                print('There was an error connecting to %s' % address)
                return
            self.set_socket(self.sock)

    def _getData(self):
        ret = self.data
        self.data = ''
        return ret
    
    def found_terminator(self):
        data = self._getData()
        
        words = data.split(" ")
        
        if words[0] == "PING":
            # The server pinged us, we need to pong or we'll be disconnected
            # (make sure to send whatever follows the PING, in case they send a random hash)
            self.sendLine("PONG " + data[5:])
            return 
        
        # Takes the ':Nick!Ident@Host' chunk and assigns (nick,ident,host) to user
        user = None
        if words[0].find("!") != -1:
            user = User(words[0][words[0].find(":") + 1:words[0].find('!')])
            user.ident = words[0][words[0].find('!') + 1:words[0].find('@')]
            user.host = words[0][words[0].find('@') + 1:]
        else:
            #it's our nick, in the format ":NickName"
            user = User(words[0][words[0].find(":") + 1:])
            
        if words[1] == "433":
            #There was a nick collision
            self.config["nickname"] = self.config["nickname"] + "_"
            self.sendLine("NICK "+self.config["nickname"])
        
        elif words[1] == "422" or words[1] == "376":
            # We successfully logged in, do post-login procedures
            if self.config["ns_name"] != "" and self.config["ns_pwd"] != "":
                self.doMessage("NickServ", "IDENTIFY "+self.config["ns_name"]+" "+self.config["ns_pwd"])
            
            self.pluginManager.event("connect", self.config["network"])
            
            for chan in [k for k in self.config["channels"].split(",") if k != '']:
                self.doJoin(chan)
        
        elif words[1] == "353":
            #user list, if it's large, we only got part of it.
            words = [k for k in data.replace(" =", "").split(" ") if k != '']
            if self.userlist.get(words[3], None) is None:
                self.userlist[words[3]] = []
            
            self.userlist[words[3]] += [u.strip(":") for u in words[3:]]
        
        elif words[1] == "PRIVMSG":
            # We got a message
            channel = (words[2] if words[2] != self.config["nickname"] else user)
            message = data[data.find(":", data.find(channel[(channel.find("-") == -1 and 1 or channel.find("-")):]))+1:]
            if message.find("\x01ACTION") == 0:
                # String was found, it's an action
                self.pluginManager.event("action", channel, user, message[8:-1])
            elif message.find(self.config["comchar"]) == 0:
                # String was found, it's a command!
                args = message.split(" ")
                cmd = args[0][len(self.config["comchar"]):]
                args.pop(0)
                self.pluginManager.event("message", channel, user, message)
                self.pluginManager.event("command", channel, user, cmd.lower(), args)
            else:
                if re.match("^"+self.config["nickname"]+"[^A-Za-z0-9] .+$", message, flags=re.IGNORECASE):
                    #Someone addressed us, it's a command (probably)
                    args = message.split(" ")
                    cmd = args[1]
                    args = args[2:]
                    self.pluginManager.event("message", channel, user, message)
                    self.pluginManager.event("command", channel, user, cmd.lower(), args)
                else:
                    #Nothing was found, it has to be a message
                    self.pluginManager.event("message", channel, user, message)
        
        elif words[1] == "JOIN":
            # Someone joined a channel that we're in
            if user != self.config["nickname"]:
                self.pluginManager.event("join", words[2].strip(":"), user)
                self.userlist[words[2].strip(":")].append(user)
        
        elif words[1] == "PART":
            if user != self.config["nickname"]:
                # Someone parted a channel we're in
                if user in self.userlist[words[2].strip(":")]:
                    self.userlist[words[2]].remove(user)
                self.pluginManager.event("part", words[2].strip(":"), user, " ".join(words[3:]))
        
        elif words[1] == "QUIT":
            # Someone quit the server
            for k in self.userlist.keys():
                if user in self.userlist[k]:
                    self.userlist[k].remove(user)
            self.pluginManager.event("quit", user, " ".join(words[2:])[1:]) 
        
        elif words[1] == "NICK":
            # Someone changed their nickname 
            for k in self.userlist.keys():
                if user in self.userlist[k]:
                    self.userlist[k].remove(user)
                    self.userlist[k].append(words[2].strip(":"))
            self.pluginManager.event("nick", user, words[2].strip(":"))
        
        elif words[1] == "MODE":
            # Someone set a mode
            try:
                self.pluginManager.event("mode", words[2], user, words[3], words[4])
            except IndexError: # words[4] is not valid, it wasn't set on a user
                self.pluginManager.event("mode", words[2], user, words[3], "")
        
        elif words[1] == "KICK":
            #someone kicked someone
            if self.userlist.get(words[2].lower(), None) is not None:
                self.userlist[words[2].lower()].remove(words[3])
                
            self.pluginManager.event("kick", words[2], user, words[3], words[4][1:])
        
        elif words[1] == "NOTICE":
            self.pluginManager.event("notice", user, words[2], " ".join(words[3:]).strip(":"))
            
        else:
            self.pluginManager.event(words[1], data)
            
    def collect_incoming_data(self, data):
        self.data += data
        
    def handle_connect(self):
        self.config["ident"] = ident = self.config["ident"] if self.config["ident"] != "" else self.config["nickname"]
        self.sendLine("USER "+ident+" * * *")
        self.sendLine("NICK "+self.config["nickname"])
    
    def handle_disconnect(self):
        self.pluginManager.event("disconnect", self.config["network"])
    
    def handle_error(self):
        traceback.print_exc()
    
    def prnt(self, line):
        print("["+self.config["network"]+"] "+line)
    
    def sendLine(self, line):
        if line.endswith("\r\n"):
            self.push(line)
        else:
            self.push(line+"\r\n")
    
    def doMessage(self, channel, message):
        self.sendLine("PRIVMSG "+channel+" :"+message)
        self.pluginManager.event("message", channel, User(self.config["nickname"]), message)
        
    def doAction(self, channel, action):
        self.sendLine("PRIVMSG "+channel+" :\x01ACTION "+action+" \x01")
        self.pluginManager.event("action", channel, User(self.config["nickname"]), action)
    
    def doQuit(self, message=""):
        self.sendLine("QUIT :" + message)
        self.pluginManager.event("quit", User(self.config["nickname"]), message)
        self.close_when_done()
    
    def doNotice(self, user, message):
        self.sendLine("NOTICE "+user+" :"+message)
    
    def doNick(self, newnick):
        self.sendLine("NICK " + newnick)
        self.pluginManager.event("nick", User(self.config["nickname"]), User(newnick))
        self.config["nickname"] = newnick
    
    def doJoin(self, channel):
        self.sendLine("JOIN "+channel)
        self.pluginManager.event("join", channel, User(self.config["nickname"]))

    def doKick(self, channel, user, message=""):
        self.sendLine("KICK %s %s :%s" % (channel, user, message))
        self.pluginManager.event("kick", channel, User(self.config["nickname"]), user, message)
        
    def doPart(self, channel, message=""):
        self.sendLine("PART "+channel)
        if channel in self.userlist.keys():
            del self.userlist[channel]
        self.pluginManager.event("part", channel, User(self.config["nickname"]), message)
    
    def doMode(self, channel, mode, user=""):
        self.sendLine("MODE "+channel+" "+mode+" "+user)
        self.pluginManager.event("mode", channel, User(self.config["nickname"]), mode, user)