twitter/ircbot.py

   1
   2 import sys
   3 import time
   4 from dateutil.parser import parse
   5 from ConfigParser import ConfigParser
   6 from heapq import heappop, heappush
   7 import traceback
   8
   9 from api import Twitter, TwitterError
  10
  11 try:
  12     import irclib
  13 except:
  14     raise ImportError(
  15         "This module requires python irclib available from "
  16         + "http://python-irclib.sourceforge.net/")
  17
  18 def debug(msg):
  19     # uncomment this for debug text stuff
  20     print >> sys.stderr, msg
  21     pass
  22
  23 class SchedTask(object):
  24     def __init__(self, task, delta):
  25         self.task = task
  26         self.delta = delta
  27         self.next = time.time()
  28
  29     def __repr__(self):
  30         return "<SchedTask %s next:%i delta:%i>" %(
  31             self.task.__name__, self.next, self.delta)
  32
  33     def __cmp__(self, other):
  34         return cmp(self.next, other.next)
  35
  36     def __call__(self):
  37         return self.task()
  38
  39 class Scheduler(object):
  40     def __init__(self, tasks):
  41         self.task_heap = []
  42         for task in tasks:
  43             heappush(self.task_heap, task)
  44
  45     def next_task(self):
  46         now = time.time()
  47         task = heappop(self.task_heap)
  48         wait = task.next - now
  49         if (wait > 0):
  50             time.sleep(wait)
  51         task()
  52         task.next = now + task.delta
  53         heappush(self.task_heap, task)
  54         debug("tasks: " + str(self.task_heap))
  55
  56     def run_forever(self):
  57         try:
  58             while True:
  59                 self.next_task()
  60         except KeyboardInterrupt:
  61             pass
  62
  63 class TwitterBot(object):
  64     def __init__(self, configFilename):
  65         self.configFilename = configFilename
  66         self.config = load_config(self.configFilename)
  67         self.irc = irclib.IRC()
  68         self.irc.add_global_handler('privmsg', self.handle_privmsg)
  69         self.ircServer = self.irc.server()
  70         self.twitter = Twitter(
  71             self.config.get('twitter', 'email'),
  72             self.config.get('twitter', 'password'))
  73         self.sched = Scheduler(
  74             (SchedTask(self.process_events, 1),
  75              SchedTask(self.check_statuses, 60)))
  76         self.lastUpdate = time.gmtime()
  77
  78     def check_statuses(self):
  79         debug("In check_statuses")
  80         try:
  81             updates = self.twitter.statuses.friends_timeline()
  82         except Exception, e:
  83             print >> sys.stderr, "Exception while querying twitter:"
  84             traceback.print_exc(file=sys.stderr)
  85             return
  86
  87         for update in updates:
  88             crt = parse(update['created_at']).utctimetuple()
  89             if (crt > self.lastUpdate):
  90                 self.privmsg_channel(
  91                     "=^_^= %s %s" %(
  92                         update['user']['screen_name'],
  93                         update['text']))
  94                 self.lastUpdate = crt
  95             else:
  96                 break
  97
  98     def process_events(self):
  99         debug("In process_events")
 100         self.irc.process_once()
 101
 102     def handle_privmsg(self, conn, evt):
 103         debug('got privmsg')
 104         args = evt.arguments()[0].split(' ')
 105         try:
 106             if (not args):
 107                 return
 108             if (args[0] == 'follow' and args[1:]):
 109                 self.follow(conn, evt, args[1])
 110             elif (args[0] == 'unfollow' and args[1:]):
 111                 self.unfollow(conn, evt, args[1])
 112             else:
 113                 conn.privmsg(
 114                     evt.source().split('!')[0],
 115                     "=^_^= Hi! I'm Twitterbot! you can (follow "
 116                     + "<twitter_name>) to make me follow a user or "
 117                     + "(unfollow <twitter_name>) to make me stop.")
 118         except Exception:
 119             traceback.print_exc(file=sys.stderr)
 120
 121     def privmsg_channel(self, msg):
 122         return self.ircServer.privmsg(
 123             self.config.get('irc', 'channel'), msg)
 124
 125     def follow(self, conn, evt, name):
 126         userNick = evt.source().split('!')[0]
 127         friends = [x['name'] for x in self.twitter.statuses.friends()]
 128         debug("Current friends: %s" %(friends))
 129         if (name in friends):
 130             conn.privmsg(
 131                 userNick,
 132                 "=O_o= I'm already following %s." %(name))
 133         else:
 134             try:
 135                 self.twitter.friendships.create(id=name)
 136             except TwitterError:
 137                 conn.privmsg(
 138                     userNick,
 139                     "=O_o= I can't follow that user. Are you sure the name is correct?")
 140                 return
 141             conn.privmsg(
 142                 userNick,
 143                 "=^_^= Okay! I'm now following %s." %(name))
 144             self.privmsg_channel(
 145                 "=o_o= %s has asked me to start following %s" %(
 146                     userNick, name))
 147
 148     def unfollow(self, conn, evt, name):
 149         userNick = evt.source().split('!')[0]
 150         friends = [x['name'] for x in self.twitter.statuses.friends()]
 151         debug("Current friends: %s" %(friends))
 152         if (name not in friends):
 153             conn.privmsg(
 154                 userNick,
 155                 "=O_o= I'm not following %s." %(name))
 156         else:
 157             self.twitter.friendships.destroy(id=name)
 158             conn.privmsg(
 159                 userNick,
 160                 "=^_^= Okay! I've stopped following %s." %(name))
 161             self.privmsg_channel(
 162                 "=o_o= %s has asked me to stop following %s" %(
 163                     userNick, name))
 164
 165     def run(self):
 166         self.ircServer.connect(
 167             self.config.get('irc', 'server'),
 168             self.config.getint('irc', 'port'),
 169             self.config.get('irc', 'nick'))
 170         self.ircServer.join(self.config.get('irc', 'channel'))
 171         try:
 172             self.sched.run_forever()
 173         except KeyboardInterrupt:
 174             pass
 175
 176 def load_config(filename):
 177     defaults = dict(server=dict(port=6667, nick="twitterbot"))
 178     cp = ConfigParser(defaults)
 179     cp.read((filename,))
 180     return cp
 181
 182 def main():
 183     configFilename = "twitterbot.ini"
 184     if (sys.argv[1:]):
 185         configFilename = sys.argv[1]
 186     bot = TwitterBot(configFilename)
 187     return bot.run()