lolbot.py 10.7 KB
Newer Older
Jonathan Harker's avatar
Jonathan Harker committed
1 2 3 4 5 6 7 8 9
#!/usr/bin/env python
#
# LolBot
#
# Code originally based on example bot and irc-bot class from
# Joel Rosdahl <joel@rosdahl.net>, author of included python-irclib.

"""
Useful bot for folks stuck behind censor walls at work
Jonathan Harker's avatar
Jonathan Harker committed
10
Logs a channel and collects URLs for later.
Jonathan Harker's avatar
Jonathan Harker committed
11 12
"""

13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28
try:
    import sys, string, random, time
    from ircbot import SingleServerIRCBot, OutputManager
    from irclib import nm_to_n, nm_to_h, irc_lower
    import os

    from datetime import datetime
    from mechanize import Browser

    import getopt
    from sqlalchemy import MetaData, Table, Column, String, Text, Integer, DateTime, engine_from_config
    from sqlalchemy.orm import sessionmaker
    from sqlalchemy.ext.declarative import declarative_base
except ImportError:
    print "Some modules could not be loaded: Lolbot relies on Mechanize and SQLAlchemy.\n"
    sys.exit
29

Jonathan Harker's avatar
Jonathan Harker committed
30 31 32 33 34 35 36 37 38 39 40 41 42 43
# Exclamations - wrong input
exclamations = [
    "Zing!",
    "Burns!",
    "Tard!",
    "Lol.",
    "Crazy!",
    "WTF?",
]

# Ponderings
ponderings = [
    "Hi, can I have a medium lamb roast, with just potatoes.",
    "Can I slurp on your Big Cock?",
Jonathan Harker's avatar
Jonathan Harker committed
44 45
    "Quentin Tarantino is so awesome I want to have his babies.",
    "No it's a week night 8pm is past my bedtime.",
Jonathan Harker's avatar
Jonathan Harker committed
46 47
]

48 49 50 51 52 53 54 55 56 57 58 59 60 61
SqlBase = declarative_base()

class Log(SqlBase):
  """
  This class represents an event in the log table and inherits from a SQLAlchemy
  convenience ORM class.
  """

  __tablename__ = "log"

  id        = Column(Integer, primary_key=True)
  timestamp = Column(DateTime)
  nickname  = Column(String(20))
  text      = Column(Text)
Jonathan Harker's avatar
Jonathan Harker committed
62

63 64 65 66
  def __init__(self, nickname, text, timestamp=None):
    if timestamp is None:
      timestamp = datetime.now()
    self.timestamp = timestamp
Jonathan Harker's avatar
Jonathan Harker committed
67
    self.nickname = nickname
68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83
    self.text = text

  def __repr__(self):
    return "(%s) %s: %s" % (self.timestamp.strftime("%Y-%m-%d %H:%M:%S"), self.nickname, self.text)

class Url(SqlBase):
  """
  This class represents a saved URL and inherits from a SQLAlchemy convenience
  ORM class.
  """

  __tablename__ = "url"

  id        = Column(Integer, primary_key=True)
  timestamp = Column(DateTime)
  nickname  = Column(String(20))
Jonathan Harker's avatar
Jonathan Harker committed
84
  url       = Column(String(200), unique=True)
85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
  title     = Column(Text)

  def __init__(self, nickname, url, title=None, timestamp=None):
    if timestamp is None:
      timestamp = datetime.now()
    self.timestamp = timestamp
    self.nickname = nickname
    self.url = url
    self.title = title

    # populate the title from the URL if not given.
    if title is None:
      try:
        br = Browser()
        br.open(self.url)
        self.title = br.title()
      except Exception as ex:
        self.title = ''

  def __repr__(self):
    if not self.title:
      return "%s: %s" % (self.nickname, self.url)
    else:
      return "%s: %s - %s" % (self.nickname, self.url, self.title)

class LolBot(SingleServerIRCBot):
  """
  The Lolbot itself.
  """

  def __init__(self, config_path=''):
    self.get_config(config_path)
    SingleServerIRCBot.__init__(self, [(self.server, self.port)], self.nickname, self.nickname)

    # connect to the database
    self.dbengine = engine_from_config(self.config, prefix="db.")
    SqlBase.metadata.bind = self.dbengine
    SqlBase.metadata.create_all()
    self.get_session = sessionmaker(bind=self.dbengine)
Jonathan Harker's avatar
Jonathan Harker committed
124 125 126

    self.helptext = "Adds URLs to a list. Commands: list - prints a bunch of URLs; clear - clears the list; lol - say something funny; <url> - adds the URL to the list; help - this message."

127
    self.queue = OutputManager(self.connection)
Jonathan Harker's avatar
Jonathan Harker committed
128 129 130
    self.queue.start()
    self.start()

131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175
  def get_config(self, config_path):
    """
    This method loads configuration options from a lolbot.conf file. The file
    should look like this:

       irc.server = irc.freenode.net
       irc.port = 6667
       irc.channel = #lolbottest
       irc.nickname = lolbot
       db.url = sqlite:///lolbot.db

    """

    if not config_path:
      config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "lolbot.conf")
    if not os.path.exists(config_path):
      print "Error: configuration file not found. By default lolbot will look for a lolbot.conf file in the same directory as the lolbot script, but you can override this by specifying a path on the command line with the --config option."
      usage()
      sys.exit(1)

    # open the configuration file and grab all param=value declarations.
    config = {}
    with open(config_path) as f:
      for line in f:
        # skip comments
        if line.strip().startswith("#") or line.strip() == "":
          continue

        # collect up param = value
        try:
          (param, value) = line.strip().split("=", 1)
          if param.strip() != "":
            config[param.strip()] = value.strip()
        except ValueError:
          continue

    # validate IRC host
    if "irc.server" not in config.keys():
      print "Error: the IRC server was not specified. Use --help for more information."
      sys.exit(1)
    self.server = config["irc.server"]

    # validate IRC port
    if "irc.port" not in config.keys():
      config["irc.port"] = "6667"
176
    try:
177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193
      self.port = int(config["irc.port"])
    except ValueError:
      print "Error: the IRC port must be an integer. If not specified, lolbot will use the default IRC port value 6667. Use --help for more information."
      sys.exit(1)

    # validate IRC channel
    if "irc.channel" not in config.keys() or not config["irc.channel"].startswith("#"):
      print "Error: the IRC channel is not specified or incorrect. It must begin with a # - e.g. #mychatchannel. Use --help for more information."
      sys.exit(1)
    self.channel = config["irc.channel"]

    # validate bot nickname
    if "irc.nickname" not in config.keys():
      config["irc.nickname"] = "lolbot"
    self.nickname = config["irc.nickname"]

    self.config = config
194

Jonathan Harker's avatar
Jonathan Harker committed
195 196 197
  def now(self):
    return datetime.today().strftime("%Y-%m-%d %H:%M:%S")

198
  def save_url(self, nickname, url):
199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216
    try:
      db = self.get_session()
      if not db.query(Url).filter(Url.url == url).count():
        theurl = Url(nickname, url)
        db.add(theurl)
        db.commit()
      else:
        try:
          theurl = db.query(Url).filter(Url.url == url).one()
          print theurl
          return theurl.title
        except MultipleResultsFound, ex:
          print ex  #wtf
        except NoResultsFound, ex:
          print ex  #wtf
    except Exception, ex:
      print "Exception caught saving URL: %s" % ex
      return ""
217 218

  def log_event(self, nick, text):
219 220 221 222 223 224 225 226
    try:
      entry = Log(nick, text)
      db = self.get_session()
      db.add(entry)
      db.commit()
      print entry
    except Exception, ex:
      print "Exception caught logging event: %s" % ex
Jonathan Harker's avatar
Jonathan Harker committed
227

Jonathan Harker's avatar
Jonathan Harker committed
228 229 230 231 232 233 234 235 236 237 238 239 240 241 242
  def on_nicknameinuse(self, connection, event):
    self.nickname = connection.get_nickname() + "_"
    connection.nick(self.nickname)

  def on_welcome(self, connection, event):
    connection.join(self.channel)

  def on_privmsg(self, connection, event):
    "Deal with a /msg private message."
    from_nick = nm_to_n(event.source())
    self.do_command(event, event.arguments()[0], from_nick)

  def on_pubmsg(self, connection, event):
    "Deal with a public message in a channel."

243
    # log it
Jonathan Harker's avatar
Jonathan Harker committed
244
    from_nick = nm_to_n(event.source())
245 246
    self.log_event(from_nick, event.arguments()[0])

Jonathan Harker's avatar
Jonathan Harker committed
247 248 249
    args = string.split(event.arguments()[0], ":", 1)
    if len(args) > 1 and irc_lower(args[0]) == irc_lower(self.nickname):
      self.do_command(event, string.strip(args[1]), from_nick)
250 251 252 253 254
    else:
      # parse it for links, add URLs to the list
      words = event.arguments()[0].split(" ")
      for w in words:
        if w.startswith('http://') or w.startswith('https://'):
255
          title = self.save_url(from_nick, w)
256
          self.say_public(title)
Jonathan Harker's avatar
Jonathan Harker committed
257 258 259 260

  def say_public(self, text):
    "Print TEXT into public channel, for all to see."
    self.queue.send(text, self.channel)
261
    self.log_event(self.nickname, text)
Jonathan Harker's avatar
Jonathan Harker committed
262 263 264 265 266 267 268 269 270 271 272 273 274

  def say_private(self, nick, text):
    "Send private message of TEXT to NICK."
    self.queue.send(text, nick)

  def reply(self, text, to_private=None):
    "Send TEXT to either public channel or TO_PRIVATE nick (if defined)."
    if to_private is not None:
      self.say_private(to_private, text)
    else:
      self.say_public(text)

  def ponder(self):
Jonathan Harker's avatar
Jonathan Harker committed
275
    "Return a random pondering."
Jonathan Harker's avatar
Jonathan Harker committed
276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292
    return random.choice(ponderings)

  def exclaim(self):
    "Return a random exclamation string."
    return random.choice(exclamations)

  def do_command(self, event, cmd, from_private):
    """
    This is the function called whenever someone sends a public or
    private message addressed to the bot. (e.g. "bot: blah").
    """

    if event.eventtype() == "pubmsg":
      target = None
    else:
      target = from_private.strip()

293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319
    try:
      if cmd == 'help':
        self.reply(self.helptext, target)

      elif cmd == 'lol':
        self.reply(self.ponder(), target)

      elif cmd == 'urls' or cmd == 'list':
        db = self.get_session()
        for url in db.query(Url).order_by(Url.timestamp):
          line = "%s %s" % (url.url, url.title)
          self.reply(line, target)
          time.sleep(1)

      elif cmd.startswith('http:') or cmd.startswith('https:'):
        title = self.save_url(from_private, cmd)
        if title == '':
          self.reply('URL added.', target)
        if title != '':
          self.reply('URL added: %s' % title, target)

      else:
        self.reply(self.exclaim(), target)

    except Exception, ex:
      print "Exception caught processing command: %s" % ex
      print "   command was '%s' from %s" % (cmd, target)
Jonathan Harker's avatar
Jonathan Harker committed
320 321


322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352
def usage():
  print """Run a lolbot.

  -h, --help
         This message.

  -c, --config=<filename>
         Specify a configuration file. Defaults to lolbot.conf in the same
         directory as the script.

Configuration:

  irc.server = <host>
     The IRC server, e.g. irc.freenode.net

  irc.port = <port>
     The IRC server port. Default: 6667

  irc.channel = <channel>
     The chat channel to join, e.g. #trainspotting

  irc.nickname = <nickname>
     The nickname for your lolbot. Default: "lolbot"

  db.url = <url>
     The URL to your lolbot database. This can be any valid URL accepted
     by SQL Alchemy. If not specified, lolbot will attempt to use a
     SQLite database called lolbot.db in the same directory as the
     script.
"""

Jonathan Harker's avatar
Jonathan Harker committed
353 354
if __name__ == "__main__":
  try:
355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370
    (options, args) = getopt.getopt(sys.argv[1:], "hc:", ["help", "config=", ])
  except getopt.GetoptError, err:
    print str(err)
    usage()
    sys.exit(2)

  config_path = ""
  for option,value in options:
    if option in ("-h", "--help"):
      usage()
      sys.exit(2)
    if option in ("-c", "--config"):
      config_path = value
  
  try:
    LolBot(config_path).start()
Jonathan Harker's avatar
Jonathan Harker committed
371 372 373
  except KeyboardInterrupt:
    print "Shutting down."