pvl/irker/irc.py
author Tero Marttila <terom@paivola.fi>
Sat, 12 Jan 2013 23:22:23 +0200
changeset 110 af87b706e4a3
parent 108 d2c1485af725
child 127 f143171884f9
permissions -rw-r--r--
pvl.irker: implement --irc-username option; improve handling of register errors
"""
    IRC client, dispatching irker messages.
"""

from twisted.internet import reactor, interfaces, protocol, defer, error
from twisted.words.protocols import irc

from twisted.internet import endpoints

from twisted.python import log

PORT = 6667

def url2endpoint (reactor, url) :
    """
        Turn given urlparse URL into an endpoint.

        Raises KeyError on unknown scheme.
    """

    SCHEMES = {
            'irc':      lambda : endpoints.TCP4ClientEndpoint(reactor, url.hostname, url.port or PORT),
    }
    
    return SCHEMES[url.scheme]()

def normalize (name) :
    """
        Normalize a channel/nickname for comparisons in IRC.
    """

    return name.lower()

class IRCError (Exception) :
    """
        A handled protocol error.
    """

    pass

class IRCChannel (object) :
    """
        A joined channel on an IRC server.
    """

    ENCODING = 'utf-8'

    def __init__ (self, client, channel, encoding=ENCODING) :
        self.client = client
        self.channel = channel

        self.encoding = encoding

    def privmsg (self, *msgs) :
        for msg in msgs :
            # XXX: encode
            self.client.msg(self.channel, msg.encode(self.encoding))

    def notice (self, *msgs) :
        for msg in msgs :
            self.client.notice(self.channel, msg.encode(self.encoding))

    def errback (self, failure) :
        """
            Fail any pending requests.
        """

        log.msg('IRCChannel.errback', self, failure)

    def __str__ (self) :
        return self.client.url(self.channel)

class IRCClient (irc.IRCClient) :
    """
        A connection to an IRC server with a specific, requested nickname.

        Joins to channels.
    """

    performLogin = False

    def __init__ (self, factory) :
        self.factory = factory
        
        self.nickname = None
        self.hostname = None

        self._registering = None
        self._channels = { }

        # TODO: smarter/configurable queueing?
        self.lineRate = 1.0

    def connectionMade (self) :
        self.hostname = self.transport.getPeer().host
        self.transport.logPrefix = self.logPrefix

        log.msg("connectionMade", self, self.transport)
        irc.IRCClient.connectionMade(self)

    def sendLine (self, line) :
        irc.IRCClient.sendLine(self, line)

        log.msg(">>>", line)

    def lineReceived (self, line) :
        log.msg("<<<", line)

        irc.IRCClient.lineReceived(self, line)
        
    ## Register
    def register (self, nickname, username=None, password=None) :
        """
            Register to the server, choosing a nickname based on the given nickname.

            Returns a Deferred that callbacks with our actual nickname once we have registered, or errbacks with an IRCError.
        """

        if self._registering :
            raise Exception("register: already registering")

        self.username = username
        self.password = password
        
        log.msg("register", nickname)
        irc.IRCClient.register(self, nickname)

        # defer
        d = self._registering = defer.Deferred()

        return d
    
    # irc_ERR_NICKNAMEINUSE
    # alterCollidedNick
    # irc_ERR_ERRONEUSNICKNAME

    def irc_ERR_PASSWDMISMATCH (self, prefix, params) :
        err = IRCError('ERR_PASSWDMISMATCH')
        log.err(err)
        self._registering.errback(err)

    def irc_RPL_WELCOME (self, prefix, params) :
        self.hostname = prefix
        irc.IRCClient.irc_RPL_WELCOME(self, prefix, params)

    def signedOn (self) :
        log.msg("signedOn", self.nickname)
        irc.IRCClient.signedOn(self)
        
        # defer
        d = self._registering

        if not d :
            raise Exception("signedOn: not registering?")

        self._registering = None

        d.callback(self.nickname)
    
    ## Channels
    def join (self, channel, key=None) :
        """
            Join the given channel.

            Returns a deferred that callbacks with the IRCChannel once joined, or errbacks.
        """

        irc.IRCClient.join(self, channel, key=key)

        d = self._channels[normalize(channel)] = defer.Deferred()

        return d

    # ERR_CHANNELISFULL
    # ERR_INVITEONLYCHAN
    # ERR_BANNEDFROMCHAN
    # ERR_BADCHANNELKEY
    
    def _close_channel (self, channel) :
        """
            Remove channel from our list of channels.

            TODO: purge queued messages for channel?
        """

        del self._channels[normalize(channel)]

    def left (self, channel) :
        log.msg('IRCClient.left', channel)

        self._close_channel(channel)

    def kickedFrom (self, channel, kicker, message) :
        log.msg('IRCClient.kicked', channel, kicker, message)
        
        self._close_channel(channel)

    def joined (self, channel) :
        """
            Have joined given channel.
        """
        
        lookup = normalize(channel)

        d = self._channels[lookup]
        channel = self._channels[lookup] = IRCChannel(self, channel)
        d.callback(channel)

    @defer.inlineCallbacks
    def channel (self, channel, key=None) :
        """
            Defer a joined IRCChannel.
        """

        lookup = normalize(channel)
        
        log.msg('IRCClient.channel', lookup, channel)

        if lookup not in self._channels :
            channel = yield self.join(channel, key)
        else :
            # wait or get
            yield self._channels[lookup]

        channel = self._channels[lookup]

        log.msg('IRCClient.channel', lookup, channel)
        
        defer.returnValue(channel)

    ## 
    def irc_ERR_CANNOTSENDTOCHAN (self, prefix, params) :
        nick, channel, error = params

        log.err(IRCError(channel, error))

    ## Quit
    def irc_ERROR (self, prefix, params) :
        msg, = params
        error = IRCError(None, msg)

        log.err(error)

        if self._registering :
            self._registering.errback(error)
            self._registering = None

    def connectionLost (self, reason) :
        irc.IRCClient.connectionLost(self, reason)
        log.err(reason)

        if self._registering :
            self._registering.errback(reason)
            self._registering = None
        
        # unregister channels
        for channel in self._channels :
            # errback Deferred or IRCChannel
            self._channels[channel].errback(reason)
        
        self._channels = { }

        # unregister client
        self.factory.clientLost(self)
    
    ## Logging
    def url (self, target=None) :
        """
            Format as URL.
        """

        if not self.transport : return 'IRC'
        
        # XXX: no isinstance() support
        if interfaces.ITCPTransport.providedBy(self.transport) :
            scheme = 'irc'
        else :
            # TODO: ssl?
            scheme = None
        
        peer = self.transport.getPeer()

        if peer.port == PORT :
            peer = "{irc.hostname}".format(irc=self, peer=peer)
        else :
            peer = "{irc.hostname}:{peer.port}".format(irc=self, peer=peer)

        if target :
            path = str(target)
        else :
            path = None
        
        return ''.join(part for part in (
            scheme, '://' if scheme else None,
            self.nickname, '@' if self.nickname else None,
            peer,
            '/' if path else None, path
        ) if part)

    __str__ = url
    logPrefix = url

class IRCFactory (protocol.ClientFactory) :
    """
        Manage Clients and Targets
    """

    NICKNAME = 'irker'

    def __init__ (self, nickname=NICKNAME, username=None) :
        # default nickname
        self.nickname = nickname
        self.username = username

        # (scheme, host, port, nick) -> IRCClient
        self.clients = {}
    
    def buildProtocol (self, addr) :
        return IRCClient(self)

    def clientLost (self, client) :
        """
            Given IRCClient is no more.
        """

        log.msg("IRCFactory.clientLost", client)
        
        # remove from our clients
        self.clients = dict((k, c) for k, c in self.clients.iteritems() if c != client)

    @defer.inlineCallbacks
    def connect (self, url) :
        """
            Defer a connected, registered Client for given URL.
        """
        
        endpoint = url2endpoint(reactor, url)

        log.msg('IRCFactory.connect', url, ':', endpoint)
        
        # connect
        try :
            client = yield endpoint.connect(self)
        
        except error.ConnectError as ex :
            log.err(ex, ': '.join(str(x) for x in ('IRCFactory.connect', url, endpoint)))
            raise
        
        else :
            log.msg('IRCFactory.connect', url, ':', endpoint, ':', client)
        
        # register
        try :
            nickname = yield client.register(url.username or self.nickname, 
                    username    = self.username,
                    password    = url.password, 
            )

        except Exception as ex :
            log.err("register", ex)
            raise
        
        log.msg('IRCFactory.connect', url, ':', endpoint, ':', client, ':', nickname)

        # okay!
        defer.returnValue(client)

    @defer.inlineCallbacks
    def client (self, url) :
        """
            Return IRCClient for given URL.
        """

        lookup = (url.scheme, url.hostname, url.port, url.username)

        if lookup not in self.clients :
            # deferred for connect
            connect = self.clients[lookup] = self.connect(url)
            
            try :
                # wait on deferred, and then store IRCClient
                self.clients[lookup] = yield connect

            except Exception as ex :
                # failed, remove the attempted connect
                del self.clients[lookup]
                raise

        else :
            # wait for result, if deferred
            # XXX: this yields None, since the first inlineCallbacks yielding on the deferred returns None in its callback
            yield self.clients[lookup]
        
        # connected client
        client = self.clients.get(lookup)

        if client :
            log.msg('IRCFactory.client', url, ":", client)

            defer.returnValue(client)
        else :
            log.msg('IRCFactory.client', url, ": client connect failed")
            
            # XXX: get failure from first yield's errback... except inlineCallbacks drops it and goes to callback with None <_<
            raise Exception("Client connect failed")

    @defer.inlineCallbacks
    def target (self, url) :
        """
            Return IRCChannel for given URL.
        """

        client = yield self.client(url)
        
        channel = '#' + url.path.lstrip('/')
        channel = yield client.channel(channel)
        
        log.msg('IRCFactory.target', url, ":", channel)

        defer.returnValue(channel)

    @defer.inlineCallbacks
    def privmsg (self, url, *msg) :
        """
            Dispatch given messages to given target.
        """

        target = yield self.target(url)
        
        log.msg('IRCFactory.privmsg', url, ":", target, ':', *msg)

        target.privmsg(*msg)