botbase.py

"""
A useful wrapper with logging etc. for the
twisted.words.protocols.irc.IRCClient class.
"""
import sys, time, fnmatch

from twisted.words.protocols import irc
from twisted.internet import protocol, reactor

DEBUG = 0
INFO = 1
WARNING = 2
ERROR = 3

def getHostmask(user):
    host = user.split("@")[1]
    return "*!*@"+host

def matchHostmask(user, hostmask):
    return fnmatch.fnmatch(user, hostmask)

class Logger(object):
    def __init__(self, level=INFO, stream=sys.stdout,
        timestamp="[%H:%M:%S]"):
        self.level = level
        self.stream = stream
        self.timestamp = timestamp

    def log(self, level, text):
        if level < self.level:
            return
        if self.timestamp:
            stamp = time.strftime(self.timestamp, time.localtime(time.time()))
            text = "%s %s" % (stamp, text)
        self.stream.write(text + "\r\n")
        self.stream.flush()

    def __call__(self, *args, **kwargs):
        self.log(*args, **kwargs)


class BotBase(irc.IRCClient):
    nickname = ""

    def __init__(self, factory):
        self.factory = factory
        BotBase.nickname = self.factory.nick #ugh
        self.logger = self.factory.logger

    # Prototype the following callbacks

    def connectionMade(self):
        irc.IRCClient.connectionMade(self)
        self.logger(INFO, "Connection made")

    def connectionLost(self, reason):
        irc.IRCClient.connectionLost(self, reason)
        self.logger(ERROR, "Connection lost (%s)" % reason)

    def signedOn(self):
        if isinstance(self.factory.channels, str):
            self.join(self.factory.channels)
        else:
            self.logger(INFO, "Joining channels...")
            for chan in self.factory.channels:
                self.join(chan)

    def joined(self, channel):
        self.logger(DEBUG, "Joined %s" % channel)

    # privmsg() dispatches a PRIVMSG into privateMsg
    # and channelMsg. Don't prototype.

    def privmsg(self, user, channel, msg):
        nick = user.split("!")[0]

        if channel == self.nickname:
            self.privateMsg(user, nick, msg)
        else:
            self.channelMsg(user, nick, channel, msg)

    # Higher-Level callbacks, prototype too.

    def privateMsg(self, user, nick, msg):
        self.logger(DEBUG, "-%s- %s" % (nick, msg))

    def channelMsg(self, user, nick, channel, msg):
        self.logger(DEBUG, "{%s} <%s> %s" % (channel, nick, msg))


class BotBaseFactory(protocol.ClientFactory):
    def __init__(self, protocol, nick, channels, args, logger=Logger()):
        self.protocol = protocol
        self.nick = nick
        self.channels = channels
        self.logger = logger
        self.reconnect = True
        self.args = args
        logger(DEBUG, "BotBaseFactory initialized")

    def buildProtocol(self, addr):
        self.logger(DEBUG, "Building Protocol...")
        p = self.protocol(self, *self.args)
        return p

    def clientConnectionLost(self, connector, reason):
        if self.reconnect:
            self.logger(INFO, "Reconnecting...")
            connector.connect()

    def clientConnectionFailed(self, connector, reason):
        self.logger(ERROR, "Connection failed (%s)" % reason)
        reactor.stop()