misc/swirc.py

#!/usr/bin/env python
#
# swirc.py
#
# An IRC bot to analyze the Small World Phenonmenon in the IRC world.
# Original idea: AlpT <alpt@freaknet.org>
# For more details see: http://idiki.dyne.org/wiki/IRC_Small_World
#
# Copyright (C) 2007  Daniele Tricoli aka Eriol <eriol@mornie.org>
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
#
# --
# Thanks to AlpT and KatolaZ for sharing this idea and coding with me at JD06!
# Thanks to MancaUSoft for helping with the absurd problems' debug ;)
#
# python-irclib sucks, long life to twisted ;)

import sys
import pydot
import time

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

import logging as log
log.basicConfig(level=log.DEBUG,
                format='[%(asctime)s] %(levelname)s %(message)s',
                datefmt='%d/%m/%Y %H:%M:%S %Z',
                filename='swirc.log',
                filemode='w')
console = log.StreamHandler()
console.setLevel(log.INFO)
log.getLogger().addHandler(console)

BOT_NICKNAME = 'swirc'
MAX_CURRENT_CHANNEL = 10

class Graph(object):

    def __init__(self, sdata=None):
        self.data = {}

        if sdata is not None:
            self.data.update(sdata)

    def addNode(self, node, reachable):
        try:
            self.data[node]
        except KeyError:
            self.data[node] = reachable
        else:
            for e in reachable:
                if e not in self.data[node]:
                    self.data[node] += reachable

    def __len__(self):
        return len(self.data)

class MakeGraph(object):

    def __init__(self):
        self.graph = Graph()

    def addNick(self, nickslist):

        for nick in nickslist:
            self.graph.addNode(nick,
                               [node for node in nickslist if node != nick])

    def _cleanEdges(self):
        edges = [(nick, x) for nick in self.graph.data
                             for x in self.graph.data[nick]]

        cleanedEdges = list(edges)

        for x in edges:
            if (x[1], x[0]) in cleanedEdges:
                del cleanedEdges[cleanedEdges.index(x)]

        return cleanedEdges

    def saveDot(self, fname):
        gr = pydot.graph_from_edges(self._cleanEdges())
        gr.write_dot(fname, prog='dot')

mgi = MakeGraph() #ugly hack to make the dot file of the graph after CTRL + C

class SwircBot(irc.IRCClient):

    nickname = BOT_NICKNAME

    def __init__(self):
        self.knownchannels = []
        self.inchannels = []
        self.knownnicks = []
        self.cnicks = []
        self.vchannels = [] # Channels we have to visit
        self.wcount = 0
        self.j = True

        self.joinNextChannel()

    def joinNextChannel(self):
        print self.wcount
        if not self.wcount and self.j:
            try:
                chan = self.vchannels.pop(0)
                self.j = False
                self.join(chan)
            except IndexError:
                pass

        reactor.callLater(5, self.joinNextChannel)

    def connectionMade(self):
        irc.IRCClient.connectionMade(self)
        self.graphmaker = mgi

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

    def signedOn(self):
        log.info('Signed On:')
        self.join(self.factory.channel)

    def joined(self, channel):
        log.info('Joined: %s', channel)
        self.j = True
        self.cnicks = []
        self.inchannels.append(channel)
        if len(self.inchannels) > MAX_CURRENT_CHANNEL:
            leaveChannel = self.inchannels.pop(0)
            self.knownchannels.append(leaveChannel)
            self.leave(leaveChannel, 'Bye')
            log.info('Leaved channel: %s', leaveChannel)

    def whois(self, nick):
        self.sendLine('WHOIS %s' % (nick,))

    def irc_RPL_NAMREPLY(self, prefix, params):
        for nick in params[3].split():
            if nick == BOT_NICKNAME:
                continue
            if nick[0] in '@+%':
                nick = nick[1:]

            self.cnicks.append(nick)

        self.graphmaker.addNick(self.cnicks)

    def irc_RPL_ENDOFNAMES(self, prefix, params):
        for nick in self.cnicks:
            if nick not in self.knownnicks:
                self.whois(nick)
                self.wcount += 1
                time.sleep(0.2)

    def irc_RPL_WHOISCHANNELS(self, prefix, params):
        for channel in params[2].split():
                if channel[0] in '@+%':
                    channel = channel[1:]

                if (channel not in self.knownchannels and
                    channel not in self.inchannels and
                    channel not in self.vchannels):
                    self.vchannels.append(channel)

    def irc_RPL_ENDOFWHOIS(self, prefix, params):
        nick = params[1]
        if nick not in self.knownnicks:
            self.knownnicks.append(nick)
        time.sleep(0.2)

        self.wcount -= 1

class SwircBotFactory(protocol.ClientFactory):

    protocol = SwircBot

    def __init__(self, channel):
        self.channel = channel

    def clientConnectionLost(self, connector, reason):
        log.critical('Connection lost: %s', reason)
        connector.connect()
        #reactor.stop()

    def clientConnectionFailed(self, connector, reason):
        log.critical('Connection failed: %s', reason)
        reactor.stop()

if __name__ == '__main__':

    f = SwircBotFactory(sys.argv[1])

    reactor.connectTCP('irc.azzurra.org', 6667, f)

    reactor.run()

    mgi.saveDot('swirc.dot')