lolbot.py 11.8 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
    self.helptext = "Keeps a list of URLs. Commands: list [n|x-y] - prints the last 10 URLs (or n URLs, or x through y); clear - clears the list; lol - say something funny; <url> - adds the URL to the list; help - this message."
Jonathan Harker's avatar
Jonathan Harker committed
126

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
    title = False
200 201 202 203 204 205 206
    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:
207 208 209
        theurl = db.query(Url).filter(Url.url == url).one()
      print theurl
      title = theurl.title
210 211
    except Exception, ex:
      print "Exception caught saving URL: %s" % ex
212
    return title
213 214

  def log_event(self, nick, text):
215 216 217 218 219 220 221 222
    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
223

Jonathan Harker's avatar
Jonathan Harker committed
224 225 226 227 228 229 230 231 232 233 234 235 236 237 238
  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."

239
    # log it
Jonathan Harker's avatar
Jonathan Harker committed
240
    from_nick = nm_to_n(event.source())
241 242
    self.log_event(from_nick, event.arguments()[0])

Jonathan Harker's avatar
Jonathan Harker committed
243 244 245
    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)
246 247 248 249 250
    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://'):
251
          title = self.save_url(from_nick, w)
252 253 254 255
          if title == False:
            self.say_public("Sorry, I'm useless at UTF-8.")
          else:
            self.say_public("URL added. %s" % title)
Jonathan Harker's avatar
Jonathan Harker committed
256 257 258 259

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

  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
274
    "Return a random pondering."
Jonathan Harker's avatar
Jonathan Harker committed
275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291
    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()

292 293 294 295 296 297 298 299 300
    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()
301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331
        for url in db.query(Url).order_by(Url.timestamp.desc())[:10]:
          line = "%s %s" % (url.url, url.title)
          self.reply(line, target)
          time.sleep(1)

      elif cmd.startswith('urls ') or cmd.startswith('list '):
        db = self.get_session()
        (listcmd, n) = cmd.split(" ", 1)
        n = n.strip()
        if n == "all":
          rows = db.query(Url).order_by(Url.timestamp.desc())
        elif n.find("-") > 0:
          (x, y) = n.split("-", 1)
          try:
            x = abs(int(x))
            y = abs(int(y))
            if y < x:
              x, y = y, x
          except ValueError, ex:
            self.reply("Give me a number or a range of numbers, e.g. list 5 or list 11-20", target)
            raise ex
          rows = db.query(Url).order_by(Url.timestamp.desc())[x-1:y]
        else:
          try:
            n = abs(int(n))
          except ValueError, ex:
            self.reply("Give me a number or a range of numbers, e.g. list 5 or list 11-20", target)
            raise ex
          rows = db.query(Url).order_by(Url.timestamp.desc())[:n]

        for url in rows:
332 333 334 335 336 337
          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)
338 339 340 341
        if title == False:
          self.say_public("Sorry, I'm useless at UTF-8.")
        else:
          self.reply('URL added. %s' % title, target)
342 343 344 345 346 347 348

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

    except Exception, ex:
      print "Exception caught processing command: %s" % ex
      print "   command was '%s' from %s" % (cmd, target)
349 350
      self.reply("Sorry, I didn't understand:  %s" % cmd, target)
      self.reply(self.helptext, target)
Jonathan Harker's avatar
Jonathan Harker committed
351 352


353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383
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
384 385
if __name__ == "__main__":
  try:
386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401
    (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
402 403 404
  except KeyboardInterrupt:
    print "Shutting down."