pybot/bot.py

#!/usr/bin/env python
#
# v0.7.1
# works with python 2.6.x and 2.7.x
#


import sys
import socket
import string
import os
import datetime
import time
import select
import traceback
import threading

import botbrain
from logger import Logger
import db
import confman
from event import Event
import db
import util

DEBUG = False

class Bot(threading.Thread):
  """
    bot instance. one bot gets instanciated per network, as an entirely distinct, sandboxed thread.
    handles the core IRC protocol stuff, and passing lines to defined events, which dispatch to their subscribed modules.
  """
  def __init__(self, conf=None, network=None, d=None):
    threading.Thread.__init__(self)

    self.DEBUG = d
    self.brain = None
    self.network = network
    self.OFFLINE = False
    self.CONNECTED = False
    self.JOINED = False
    self.conf = conf
    self.db = db.DB()
    self.pid = os.getpid()
    self.logger = Logger()

    self.NICK = self.conf.getNick(self.network)

    self.logger.write(Logger.INFO, "\n", self.NICK)
    self.logger.write(Logger.INFO, " initializing bot, pid " + str(os.getpid()), self.NICK)

    # arbitrary key/value store for modules
    # they should be 'namespaced' like bot.mem_store.module_name
    self.mem_store = dict()

    self.CHANNELINIT = conf.getChannels(self.network)
# this will be the socket
    self.s = None # each bot thread holds its own socket open to the network
    self.brain = botbrain.BotBrain(self.send, self) 

    self.events_list = list()

    # define events here and add them to the events_list

    all_lines = Event("1__all_lines__")
    all_lines.define(".*")
    self.events_list.append(all_lines)

    implying = Event("__implying__")
    implying.define(">")

    #command = Event("__command__")
   # this is an example of passing in a regular expression to the event definition
    #command.define("fo.bar")

    lastfm = Event("__.lastfm__")
    lastfm.define(".lastfm")

    dance = Event("__.dance__")
    dance.define("\.dance")

    #unloads = Event("__module__")
    #unloads.define("^\.module")

    pimp = Event("__pimp__")
    pimp.define("\.pimp")

    bofh = Event("__.bofh__")
    bofh.define("\.bofh")

    #youtube = Event("__youtubes__")
    #youtube.define("youtube.com[\S]+")

    weather = Event("__.weather__")
    weather.define("\.weather")

    steam = Event("__.steam__")
    steam.define("\.steam")

    part = Event("__.part__")
    part.define("part")

    tell = Event("__privmsg__")
    tell.define("PRIVMSG")

    links = Event("__urls__")
    links.define("https?://*")

  # example
  #  test = Event("__test__")
  #  test.define(msg_definition="^\.test")

    # add your defined events here
    # tell your friends
    self.events_list.append(lastfm)
    self.events_list.append(dance)
    self.events_list.append(pimp)
    #self.events_list.append(youtube)
    self.events_list.append(bofh)
    self.events_list.append(weather)
    self.events_list.append(steam)
    self.events_list.append(part)
    self.events_list.append(tell)
    self.events_list.append(links)

    self.time_event = Event("__time_test__")
    self.time_event.define(time_event=True)
    self.events_list.append(self.time_event)

    #self.events_list.append(unloads)
  # example
  #  self.events_list.append(test)

    self.load_modules()
    self.logger.write(Logger.INFO, "bot initialized.", self.NICK)

  # conditionally subscribe to events list or add event to listing
  def register_event(self, event, module):
    """
      Allows for dynamic, asynchronous event creation. To be used by modules, mostly, to define their own events in their initialization.
      Prevents multiple of the same _type of event being registered.

      Args: 
      event: an event object to be registered with the bot 
      module: calling module; ensures the calling module can be subscribed to the event if it is not already.

      Returns: 
      nothing.
    """
    for e in self.events_list:
      if e.definition == event.definition and e._type == event._type:
        # if our event is already in the listing, don't add it again, just have our module subscribe
        e.subscribe(module)
        return

    self.events_list.append(event)
    return

  # utility function for loading modules; can be called by modules themselves
  def load_modules(self, specific=None):
    """
      Run through the ${bot_dir}/modules directory, dynamically instantiating each module as it goes.

      Args:
      specific: string name of module. if it is specified, the function attempts to load the named module.

      Returns:
      1 if successful, 0 on failure. In keeping with the perverse reversal of UNIX programs and boolean values.
    """
    nonspecific = False
    found = False

    self.loaded_modules = list()

    modules_dir_list = list()
    tmp_list = list()
    
    modules_path = 'modules'
    autoload_path = 'modules/autoloads'

    # this is magic.

    import inspect
    import os, imp, json
    dir_list = os.listdir(modules_path)
    mods = {}
    autoloads = {}
    #load autoloads if it exists
    if os.path.isfile(autoload_path): 
      self.logger.write(Logger.INFO, "Found autoloads file", self.NICK)
      try:
        autoloads = json.load(open(autoload_path))
        #logging
        for k in autoloads.keys():
          self.logger.write(Logger.INFO, "Autoloads found for network " + k, self.NICK)
      except IOError:
        self.logger.write(Logger.ERROR, "Could not load autoloads file.",self.NICK)
    # create dictionary of things in the modules directory to load
    for fname in dir_list:
      name, ext = os.path.splitext(fname)
      if specific is None:
        nonspecific = True
        # ignore compiled python and __init__ files. 
        #choose to either load all .py files or, available, just ones specified in autoloads
        if self.network not in autoloads.keys(): # if autoload does not specify for this network
          if ext == '.py' and not name == '__init__': 
            f, filename, descr = imp.find_module(name, [modules_path])
            mods[name] = imp.load_module(name, f, filename, descr)
            self.logger.write(Logger.INFO, " loaded " + name + " for network " + self.network, self.NICK)
        else: # follow autoload's direction
          if ext == '.py' and not name == '__init__':
            if name == 'module':
              f, filename, descr = imp.find_module(name, [modules_path])
              mods[name] = imp.load_module(name, f, filename, descr)
              self.logger.write(Logger.INFO, " loaded " + name + " for network " + self.network, self.NICK)
            elif ('include' in autoloads[self.network] and name in autoloads[self.network]['include']) or ('exclude' in autoloads[self.network] and name not in autoloads[self.network]['exclude']):
              f, filename, descr = imp.find_module(name, [modules_path])
              mods[name] = imp.load_module(name, f, filename, descr)
              self.logger.write(Logger.INFO, " loaded " + name + " for network " + self.network, self.NICK)
      else:
        if name == specific: # we're reloading only one module
          if ext != '.pyc': # ignore compiled 
            f, filename, descr = imp.find_module(name, [modules_path])
            mods[name] = imp.load_module(name, f, filename, descr)
            found = True

    for k,v in mods.iteritems(): 
      for name in dir(v):
        obj = getattr(mods[k], name) # get the object from the namespace of 'mods'
        try:
          if inspect.isclass(obj): # it's a class definition, initialize it
            a = obj(self.events_list, self.send, self, self.say) # now we're passing in a reference to the calling bot
            if a not in self.loaded_modules: # don't add in multiple copies
              self.loaded_modules.append(a)
        except TypeError:
          pass

    if nonspecific is True or found is True:
      return 0
    else: return 1
    # end magic.
    
  def send(self, message):
    """
    Simply sends the specified message to the socket. Which should be our connected server.
    We parse our own lines, as well, in case we want an event triggered by something we say.
    If debug is True, we also print out a pretty thing to console.

    Args:
    message: string, sent directly and without manipulation (besides UTF-8ing it) to the server.
    """
    if self.OFFLINE:
      self.debug_print(util.bcolors.YELLOW + " >> " + util.bcolors.ENDC + self.getName() + ": " + message.encode('utf-8', 'ignore'))
    else:
      if self.DEBUG is True:
        self.logger.write(Logger.INFO, str(datetime.datetime.now()) + ": " + self.getName() + ": " + message.encode('utf-8', 'ignore'), self.NICK)
        self.debug_print(util.bcolors.OKGREEN + ">> " + util.bcolors.ENDC + ": " + self.getName() + ": " + message.encode('utf-8', 'ignore'))

      self.s.send(message.encode('utf-8', 'ignore'))
      target = message.split()[1]
      if target.startswith("#"):
        self.processline(':' + self.conf.getNick(self.network) + '!~' + self.conf.getNick(self.network) + '@fakehost.here ' + message.rstrip()) 

  def pong(self, response):
    """
    Keepalive heartbeat for IRC protocol. Until someone changes the IRC spec, don't modify this.
    """
    self.send('PONG ' + response + '\n')

  def processline(self, line):
    """
    Grab newline-delineated lines sent to us, and determine what to do with them. 
    This function handles our initial low-level IRC stuff, as well; if we haven't joined, it waits for the MOTD message (or message indicating there isn't one) and then issues our own JOIN calls.

    Also immediately passes off PING messages to PONG.

    Args:
    line: string. 

    """
    if self.DEBUG:
      import datetime
      self.debug_print(util.bcolors.OKBLUE + "<< "  + util.bcolors.ENDC + ": " + self.getName() + ": " + line)

    message_number = line.split()[1]

    try:
      for e in self.events_list:
        if e.matches(line):
          e.notifySubscribers(line)
      # don't bother going any further if it's a PING/PONG request
      if line.startswith("PING"):
        ping_response_line = line.split(":", 1)
        self.pong(ping_response_line[1])
      # pings we respond to directly. everything else...
      else:
# patch contributed by github.com/thekanbo
        if self.JOINED is False and (message_number == "376" or message_number == "422"): 
          # wait until we receive end of MOTD before joining, or until the server tells us the MOTD doesn't exis
          self.chan_list = self.conf.getChannels(self.network) 
          for c in self.chan_list:
            self.send('JOIN '+c+' \n')
          self.JOINED = True
        
        line_array = line.split()
        user_and_mask = line_array[0][1:]
        usr = user_and_mask.split("!")[0]
        channel = line_array[2]
        try: 
          message = line.split(":",2)[2]
          self.brain.respond(usr, channel, message)
        except IndexError:
          try:
            message = line.split(":",2)[1]
            self.brain.respond(usr, channel, message)
          except IndexError:
            print "index out of range.", line
        
    except Exception:
      print "Unexpected error:", sys.exc_info()[0]
      traceback.print_exc(file=sys.stdout)

      
  def worker(self, mock=False):
    """
    Open the socket, make the first incision^H^H connection and get us on the server. 
    Handles keeping the connection alive; if we disconnect from the server, attempts to reconnect.

    Args:
    mock: boolean. If mock is true, don't loop forever -- mock is for testing.
    """
    self.HOST = self.network
    self.NICK = self.conf.getNick(self.network)

    # we have to cast it to an int, otherwise the connection fails silently and the entire process dies
    self.PORT = int(self.conf.getPort(self.network))
    self.IDENT = 'mypy'
    self.REALNAME = 's1ash'
    self.OWNER = self.conf.getOwner(self.network) 
    
    # connect to server
    self.s = socket.socket()
    while self.CONNECTED == False:
      try:
# low level socket TCP/IP connection
        self.s.connect((self.HOST, self.PORT)) # force them into one argument
        self.CONNECTED = True
        self.logger.write(Logger.INFO, "Connected to " + self.network, self.NICK)
        if self.DEBUG:
          self.debug_print(util.bcolors.YELLOW + ">> " + util.bcolors.ENDC + "connected to " + self.network)
      except:
        self.debug_print(util.bcolors.FAIL + ">> " + util.bcolors.ENDC + "Could not connect! Retrying... ")
        time.sleep(1)
        self.worker()

    # core IRC protocol stuff
      self.s.send('NICK '+self.NICK+'\n')

      if self.DEBUG:
        self.debug_print(util.bcolors.YELLOW + ">> " + util.bcolors.ENDC + self.network + ': NICK ' + self.NICK + '\\n')

      self.s.send('USER '+self.IDENT+ ' 8 ' + ' bla : '+self.REALNAME+'\n') # yeah, don't delete this line

      if self.DEBUG:
        self.debug_print(util.bcolors.YELLOW + ">> " + util.bcolors.ENDC + self.network + ": USER " +self.IDENT+ ' 8 ' + ' bla : '+self.REALNAME+'\\n')

      time.sleep(3) # allow services to catch up

      self.s.send('PRIVMSG nickserv identify '+self.conf.getIRCPass(self.network)+'\n')  # we're registered!

      if self.DEBUG:
        self.debug_print(util.bcolors.YELLOW + ">> " + util.bcolors.ENDC + self.network + ': PRIVMSG nickserv identify '+self.conf.getIRCPass(self.network)+'\\n')

    self.s.setblocking(1)
    
    read = ""
    
    timeout = 0

#   if we're only running a test of connecting, and don't want to loop forever
    if mock:
      return
    # infinite loop to keep parsing lines
    while True:
      try:
        timeout += 1
        # if we haven't received anything for 120 seconds
        time_since = self.conf.getTimeout(self.network)
        if timeout > time_since:
          if self.DEBUG:
            print "Disconnected! Retrying... "
          self.logger.write(Logger.CRITICAL, "Disconnected!", self.NICK)
# so that we rejoin all our channels upon reconnecting to the server
          self.JOINED = False
          self.CONNECTED = False
          self.worker()

        if self.JOINED:
          self.time_event.time_notify_subscribers()
        time.sleep(1)
       # if self.CONNECTED == False:
       #   self.connect()
        ready = select.select([self.s],[],[], 1)
        if ready[0]:
          try:
            read = read + self.s.recv(1024)
          except:
            if self.DEBUG:
              print "Disconnected! Retrying... "
            self.logger.write(Logger.CRITICAL, "Disconnected!", self.NICK)
            self.JOINED = False
            self.CONNECTED = False
            self.worker()

          lines = read.split('\n')
          
          # Important: all lines from irc are terminated with '\n'. lines.pop() will get you any "to be continued"
          # line that couldn't fit in the socket buffer. It is stored and tacked on to the start of the next recv.
          read = lines.pop() 

          if len(lines) > 0:
            timeout = 0

          for line in lines:
            line = line.rstrip()
            self.processline(line)      
      except KeyboardInterrupt:
        print "keyboard interrupt caught; exiting ..."
        raise
  # end worker

  def debug_print(self, line):
    """
    Prepends incoming lines with the current timestamp and the thread's name, then spits it to stdout.
    Args:
    line: text.
    """
    print str(datetime.datetime.now()) + ": " + self.getName() + ": " + line.strip('\n')

  def commands(*command_list):
# stolen shamelessly from willie. damn, this is a good idea.
    """Decorator. Sets a command list for a callable.

    This decorator can be used to add multiple commands to one callable in a
    single line. The resulting match object will have the command as the first
    group, rest of the line, excluding leading whitespace, as the second group.
    Parameters 1 through 4, seperated by whitespace, will be groups 3-6.

    Args:
    command: A string, which can be a regular expression.

    Returns:
    A function with a new command appended to the commands
    attribute. If there is no commands attribute, it is added.

    Example:
    @command("hello"):
    If the command prefix is "\.", this would trigger on lines starting
    with ".hello".

    @commands('j', 'join')
    If the command prefix is "\.", this would trigger on lines starting
    with either ".j" or ".join".

    """
    def add_attribute(function):
      if not hasattr(function, "commands"):
        function.commands = []
      function.commands.extend(command_list)
      return function
    return add_attribute

  def depends(self, module_name):
    for m in self.loaded_modules:
      if m.__class__.__name__ == module_name:
        return m
    return None
    
  def run(self):
    """
    For implementing the parent threading.Thread class. Allows the thread the be initialized with our code.
    """
    self.worker()

  def say(self, channel, thing):
    """
    Speak, damn you!
    """
    self.brain.say(channel, thing)
# end class Bot
              
## MAIN ## ACTUAL EXECUTION STARTS HERE

if __name__ == "__main__":
  import bot
  DEBUG = False
  for i in sys.argv:
    if i == "-d":
      DEBUG = True

# duuude this is so old.
  if os.name == "posix":
    botslist = list()
    if not DEBUG:
      pid = os.fork()
      if pid == 0: # child
        print "starting bot in the background, pid " + util.bcolors.GREEN + str(os.getpid()) + util.bcolors.ENDC
        if len(sys.argv) > 1:
          CONF = sys.argv[1]
        else: CONF = "~/.pybotrc"

        cm = confman.ConfManager(CONF)
        net_list = cm.getNetworks()
        i = 0
        if cm.getNumNets() > 1:
          for c in cm.getNetworks():
            b = bot.Bot(cm, net_list[i], DEBUG)
            b.start()
            i += 1
        else:
          b = bot.Bot(cm, net_list[0], DEBUG)
          b.start()
      elif pid > 0:
        sys.exit(0)
    else: # don't background
      print "starting bot in the background, pid " + util.bcolors.GREEN + str(os.getpid()) + util.bcolors.ENDC
      if len(sys.argv) > 1 and sys.argv[1] != "-d": # the conf file must be first argument
        CONF = sys.argv[1]
        try:
          f = open(CONF)
        except IOError:
          print "Could not open conf file " + sys.argv[1]
          sys.exit(1)
      else: CONF = "~/.pybotrc"

      cm = confman.ConfManager(CONF)
      net_list = cm.getNetworks()
      i = 0
      if cm.getNumNets() > 1:
        for c in cm.getNetworks():
          try:
            b = bot.Bot(cm, net_list[i], DEBUG)
            b.daemon = True
            b.start()
            botslist.append(b)
            i += 1
            while True: time.sleep(5)
          except (KeyboardInterrupt, SystemExit):
            print "keyboard interrupt caught; exiting..."
            sys.exit(0)
      else:
        try:
          b = bot.Bot(cm, net_list[0], DEBUG)
          b.daemon = True
          b.start()
          botslist.append(b)
          while True: time.sleep(5)
        except (KeyboardInterrupt, SystemExit):
          print "keyboard interrupt caught; exiting..."
          sys.exit(0)