An extensible IRC bot, including logging to MongoDB.

irc.py

#!/usr/bin/env python
# encoding: utf-8

"""An implementation of a very simple IRC bot using the Protocol wrappers.  This is the main script.

Requires: pymongo
"""

import logging
import re

from protocol import LineProtocol

from plugins import *


log = logging.getLogger(__name__)


class IRCProtocol(LineProtocol):
    def __init__(self, host, port, nick="arkivo", user="Arkivo", name="Arkivo IRC Logger", password=None):
        super(IRCProtocol, self).__init__(host, port)
        
        self.nick = nick
        self.user = user
        self.name = name
        self.password = password
        
        self.plugins = {}
        
    def __call__(self):
        """This protocol only implements an IRC client, default to a client."""
        
        super(IRCProtocol, self).__call__(False)
    
    def listen(self, *args, **kw):
        """This protocol only implements an IRC client."""
        
        raise NotImplementedError()
    
    def write(self, line):
        """We add debug logging of all data written to the server."""
        log.info("<- %s", unicode(line).decode('utf-8', 'replace'))
        super(IRCProtocol, self).write(line.encode('utf-8'))
    
    def send(self, text=None, *args):
        if text:
            text = text.replace('\n', '').replace('\r', '')
        
        args = [arg.replace('\n', '').replace('\r', '') for arg in args]
        
        self.write(' '.join(args) + ('' if not text else (' :' + text + '\r\n')))
    
    def connect(self):
        log.info("Connecting to irc://%s:%s...", *self.address)
        super(IRCProtocol, self).connect()
    
    def connected(self):
        log.warn("Connected to irc://%s:%s.", *self.address)
        
        if self.password: 
            self.send(None, 'PASS', self.password)
        
        self.send(None, 'NICK', self.nick)
        self.send(self.name, 'USER', self.user, '+iw', self.nick)
              
        super(IRCProtocol, self).connected()
    
    def stopped(self):
        log.warn("Disconnected from irc://%s:%s.", *self.address)
    
    def register(self, plugin, *args, **kw):
        instance = plugin(*args, **kw)
        
        for syntax in instance.syntax:
            self.plugins[re.compile(syntax)] = instance
    
    def process(self, line):
        if line[0] == ':':
            line = line[1:]
        
        log.info("-> %s", line.decode('utf-8', 'replace'))
        
        for syntax, plugin in self.plugins.iteritems():
            match = syntax.search(line)
            
            if match:
                # log.debug("Matched %r: %r", plugin, match)
                if plugin(self, **match.groupdict()):
                    return True


if __name__ == '__main__':
    logging.basicConfig(level=logging.WARN)
    
    protocol = IRCProtocol("irc.freenode.net", 6667, 'mvp-bot')
    
    protocol.register(Ping)
    protocol.register(UserPing)
    protocol.register(Echo)
    protocol.register(Control)
    protocol.register(Logger, 'localhost', 'arkivo')
    protocol.register(Join, channels=['mvp-devs', 'turbogears', 'webcore'])
    
    protocol()

plugins.py

# encoding: utf-8

"""A collection of plugins for the simple IRC bot."""

import datetime


particles = u"""able about above abst accordance according accordingly across act actually added adj adopted affected affecting affects after afterwards again against all almost alone along already also although always among amongst and announce another any anybody anyhow anymore anyone anything anyway anyways anywhere apparently approximately are aren arent aren't arise around aside ask asking auth available away awfully back became because become becomes becoming been before beforehand begin beginning beginnings begins behind being believe below beside besides between beyond biol both brief briefly but came can cannot can't cause causes certain certainly come comes contain containing contains could couldnt date did didn't different does doesn't doing done don't down downwards due during each edu effect eight eighty either else elsewhere end ending enough especially et-al etc even ever every everybody everyone everything everywhere except far few fifth first five fix followed following follows for former formerly forth found four from further furthermore gave get gets getting give given gives giving goes gone got gotten had happens hardly has hasn't have haven't having hed hence her here hereafter hereby herein heres hereupon hers herself hes hid him himself his hither home how howbeit however hundred i'll i'm immediate immediately importance important indeed index information instead into invention inward isn't itd it'd it'll its itself i've just jk keep keeps kept keys know known knows largely last lately later latter latterly least less lest let lets like liked likely line little 'll look looking looks ltd made mainly make makes many may maybe mean means meantime meanwhile merely might million miss more moreover most mostly mrs much mug must myself name namely nay near nearly necessarily necessary need needs neither never nevertheless new next nine ninety nobody non none nonetheless noone nor normally nos not noted nothing now nowhere obtain obtained obviously off often okay old omitted once one ones only onto ord other others otherwise ought our ours ourselves out outside over overall owing own page pages part particular particularly past per perhaps placed please plus poorly possible possibly potentially predominantly present previously primarily probably promptly proud provides put que quickly quite ran rather readily really recent recently ref refs regarding regardless regards related relatively research respectively resulted resulting results right run said same saw say saying says sec section see seeing seem seemed seeming seems seen self selves sent seven several shall she shed she'll shes should shouldn't show showed shown showns shows significant significantly similar similarly since six slightly some somebody somehow someone somethan something sometime sometimes somewhat somewhere soon sorry specifically specified specify specifying state states still stop strongly sub substantially successfully such sufficiently suggest sup sure take taken taking tell tends than thank thanks thanx that that'll thats that've the their theirs them themselves then thence there thereafter thereby thered therefore therein there'll thereof therere theres thereto thereupon there've these they theyd they'll theyre they've think this those thou though thoughh thousand throug through throughout thru thus til tip to together too took toward towards tried tries truly try trying twice two under unfortunately unless unlike unlikely until unto up upon ups us use used useful usefully usefulness uses using usually value various 've very via viz vol vols want wants was wasn't way wed welcome we'll went were weren't we've what whatever what'll whats when whence whenever where whereafter whereas whereby wherein wheres whereupon wherever whether which while whim whither who whod whoever whole who'll whom whomever whos whose why widely willing wish with within without won't words world would wouldn't www yes yet you youd you'll your youre yours yourself yourselves you've zero what're""".split()


class Plugin(object):
    syntax = []
    
    def __call__(self):
        raise NotImplementedError()
    
    def idle(self):
        pass


class Ping(Plugin):
    syntax = [r'^PING (?P<origin>.+)$'] # [r'^\!ping(?:[ \t]+(\d+))?$']
    
    def __call__(self, connection, origin):
        # TODO: Handle ".ping".
        connection.send(None, 'PONG', origin)

class Join(Plugin):
    """!join <channel> - Have this software join an additional channel."""
    
    syntax = [
            r'^(?P<nick>)\S+ MODE \S+ :.*$',
            r'^(?P<nick>[^!]*)!?[^@]*@?\S* PRIVMSG \S+ :\!join (?P<channel>.+)$'
        ]
    
    def __init__(self, channels):
        super(Join, self).__init__()
        self.channels = channels
        self.connected = []
    
    def __call__(self, connection, nick, channel=None):
        # TODO: Deal with removals, too.  Diffing the connected and channels lists should do.
        
        if channel:
            if nick != "GothAlice":
                connection.send("I'm sorry, certain commands are restricted to authorized personnel.", "PRIVMSG", nick)
                return
            
            if channel[1:] not in self.connected:
                connection.send(channel, "JOIN")
                self.connected.append(channel[1:])
            
            return
        
        for channel in self.channels:
            if channel not in self.connected:
                connection.send('#' + channel, "JOIN")
                self.connected.append(channel)

class UserPing(Plugin):
    """!ping - Respond as soon as possible with a short response."""
    
    syntax = [r'^(?P<origin>[^!]*)!?[^@]*@?\S* PRIVMSG (?P<target>\S+) :\!ping$']
    
    def __call__(self, connection, origin, target):
        if target[0] != '#':
            # Reply directly.
            connection.send("pong", "PRIVMSG", origin)
            return
        
        # Reply in-channel.
        connection.send("%s: pong" % (origin, ), "PRIVMSG", target)


class Echo(Plugin):
    """!echo <message> - Respond as soon as possible with copy of the given message."""
    
    syntax = [r'^(?P<origin>[^!]*)!?[^@]*@?\S* PRIVMSG (?P<target>\S+) :\!echo (?P<message>.+)$']
    
    def __call__(self, connection, origin, target, message):
        if target[0] != '#':
            # Reply directly.
            connection.send(message, "PRIVMSG", origin)
            return
        
        # Reply in-channel.
        connection.send(message, "PRIVMSG", target)


class Logger(Plugin):
    syntax = [r'^(?P<origin>[^!]*)!?[^@]*@?\S* PRIVMSG (?P<target>\S+) :(?P<message>.+)$']
    
    def __init__(self, db_host, db_name):
        import pymongo
        
        host, sep, port = db_host.partition(":")
        if not port: port = None
        
        self.db = pymongo.Connection(host, port)[db_name]
    
    def __call__(self, connection, origin, target, message):
        action = False
        kind = "message"
        message = message.decode('utf-8')
        
        if message.startswith('ACTION '):
            kind = "action"
            message = message[7:]
        
        if target[0] == '#':
            target = 'channel.' + target[1:]
        
        else:
            target = 'user.' + origin
        
        collection = self.db[target]
        
        entry = dict(
                author = origin,
                text = message,
                when = datetime.datetime.utcnow(),
                kind = kind,
                keywords = list(set([i.lower().strip(' \t.,\'"()[]{}?!:;*/\\^') for i in message.split() if i.lower() not in particles and len(i) > 2]))
            )
        
        collection.insert(entry)


class Control(Plugin):
    """!quit - Shut down this software.  Administrators only."""
    
    syntax = [r'^(?P<origin>[^!]*)!?[^@]*@?\S* PRIVMSG (?P<target>\S+) :\!quit$']
    
    def __call__(self, connection, origin, target):
        if target != connection.nick:
            log.warn("Attempt by %s to shut down the robot in the %s channel.", origin, target)
            return
        
        if origin != "GothAlice":
            log.warn("Attempt by %s to shut down the robot.", origin)
            connection.send("I'm sorry, certain commands are restricted to authorized personnel.", "PRIVMSG", origin)
            return
        
        if message == '.quit':
            return True

protocol.py

# encoding: utf-8

"""A somewhat simplified client/server sockets wrapper."""

import logging
import socket

log = logging.getLogger(__name__)
CRLF = "\r\n"


class Protocol(object):
    def __init__(self, host, port):
        super(Protocol, self).__init__()
        
        self.socket = None
        self.address = (host if host is not None else '', port)
        self.running = False
        self.server = None
    
    def __call__(self, serve=True):
        try:
            if serve:
                self.listen()
            
            else:
                self.connect()
        
        except KeyboardInterrupt:
            log.info("Recieved Control+C.")
        
        except SystemExit:
            log.info("Recieved SystemExit.")
            raise
        
        except:
            log.exception("Unknown server error.")
            raise
        
        finally:
            self.stop()
    
    def _socket(self):
        host, port = self.address
        
        try:
            addr, family, kind, protocol, name, sa = ((host, port), ) + socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM, 0, socket.AI_PASSIVE)[0]
        
        except socket.gaierror:
            if ':' in host:
                addr, family, kind, protocol, name, sa = ((host, port), socket.AF_INET6, socket.SOCK_STREAM, 0, "", (host, port, 0, 0))
            
            else:
                addr, family, kind, protocol, name, sa = ((host, port), socket.AF_INET, socket.SOCK_STREAM, 0, "", (host, port))
        
        sock = socket.socket(family, kind, protocol)
        # fixes.prevent_socket_inheritance(sock)

        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

        # TODO: Allow TCP_NODELAY sockets.
        if False:
            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

        # If listening on the IPV6 any address ('::' = IN6ADDR_ANY), activate dual-stack.
        if family == socket.AF_INET6 and addr[0] in ('::', '::0', '::0.0.0.0'):
            try:
                sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)

            except (AttributeError, socket.error):
                pass
        
        return sock
    
    def listen(self, pool=5):
        if self.running:
            raise Exception("Already running.", self.running)
        
        self.socket = self._socket()
        # self.socket.settimeout(1)
        self.socket.bind(self.address)
        self.socket.listen(pool)

        self.running = True
        self.listening()
    
    def listening(self):
        """Override this to implement asynchronous sockets, etc."""
        
        while self.running:
            sock, address = self.socket.accept()
            
            connection = self.__class__(address[0], address[1])
            connection.server = self
            connection.socket = sock
            connection.running = True
            
            connection.connection()
            
            sock.shutdown(2)
            sock.close()
    
    def connection(self, sock, address):
        """Implement this in your own sub-class to perform some useful action with new clients."""
        
        raise NotImplementedError()
    
    def connect(self):
        if self.running:
            raise Exception("Already running.", self.running)

        self.socket = self._socket()
        self.socket.settimeout(10)
        self.socket.connect(self.address)
        
        self.running = True
        self.connected()
    
    def connected(self):
        raise NotImplementedError()
    
    def stop(self):
        if not self.running:
            return
        
        self.running = False
        
        if not self.socket:
            return
        
        try:
            self.socket.shutdown(2)
            self.socket.close()
        
        except:
            log.exception("Error stopping the protocol.")
        
        self.stopped()
    
    def stopped(self):
        pass



class LineProtocol(Protocol):
    def __init__(self, host, port, buffer=1024, separator=CRLF):
        super(LineProtocol, self).__init__(host, port)
        self.buffer = buffer
        self.separator = separator
    
    def write(self, line):
        self.socket.sendall(line + self.separator)
    
    def tick(self):
        log.debug("**")
    
    def connected(self):
        """If your protocol has HELO/BYE chatter, override this method."""
        
        buf = ""
        size = self.buffer
        sock = self.socket
        separator = self.separator
        
        while True:
            try:
                buf = buf + sock.recv(size)
            
            except socket.timeout:
                if self.tick():
                    return
                
                continue
            
            while True:
                line, sep, buf = buf.partition(separator)
            
                if not sep:
                    buf = line
                    break
                
                if self.process(line):
                    return
    
    connection = connected
    
    def process(self, line):
        """You can determine if we are running as a server if self.server is set."""
        raise NotImplementedError()


class EchoProtocol(LineProtocol):
    def process(self, line):
        log.info("%s", line)
    
        self.write(line)
        
        if line == 'quit':
            return True
        
        if line == 'shutdown':
            self.server.running = False
            return True


if __name__ == '__main__':
    logging.basicConfig(level=logging.DEBUG)
    
    protocol = EchoProtocol(None, 8000)
    protocol()