lolbot.py 10.1 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
"""

import sys, string, random, time
14
from ircbot import SingleServerIRCBot, OutputManager
Jonathan Harker's avatar
Jonathan Harker committed
15
from irclib import nm_to_n, nm_to_h, irc_lower
Jonathan Harker's avatar
Jonathan Harker committed
16 17 18
import os

from datetime import datetime
19
from mechanize import Browser
Jonathan Harker's avatar
Jonathan Harker committed
20

21 22 23 24 25
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

Jonathan Harker's avatar
Jonathan Harker committed
26 27 28 29 30 31 32 33 34 35 36 37 38 39 40
# 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?",
    "Your Mum likes it two in the pink one in the stink.",
Jonathan Harker's avatar
Jonathan Harker committed
41 42
    "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
43 44
]

45 46 47 48 49 50 51 52 53 54 55 56 57 58
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
59

60 61 62 63
  def __init__(self, nickname, text, timestamp=None):
    if timestamp is None:
      timestamp = datetime.now()
    self.timestamp = timestamp
Jonathan Harker's avatar
Jonathan Harker committed
64
    self.nickname = nickname
65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80
    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
81
  url       = Column(String(200), unique=True)
82 83 84 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
  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
121 122 123

    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."

124
    self.queue = OutputManager(self.connection)
Jonathan Harker's avatar
Jonathan Harker committed
125 126 127
    self.queue.start()
    self.start()

128 129 130 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
  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"
173
    try:
174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190
      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
191

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

195 196
  def save_url(self, nickname, url):
    db = self.get_session()
Jonathan Harker's avatar
Jonathan Harker committed
197 198 199 200 201 202 203 204 205 206 207
    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()
      except MultipleResultsFound, ex:
        print ex  #wtf
      except NoResultsFound, ex:
        print ex  #wtf
208 209 210 211 212 213 214 215 216
    print theurl
    return theurl.title

  def log_event(self, nick, text):
    entry = Log(nick, text)
    db = self.get_session()
    db.add(entry)
    db.commit()
    print entry
Jonathan Harker's avatar
Jonathan Harker committed
217

Jonathan Harker's avatar
Jonathan Harker committed
218 219 220 221 222 223 224 225 226 227 228 229 230 231 232
  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."

233
    # log it
Jonathan Harker's avatar
Jonathan Harker committed
234
    from_nick = nm_to_n(event.source())
235 236
    self.log_event(from_nick, event.arguments()[0])

Jonathan Harker's avatar
Jonathan Harker committed
237 238 239
    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)
240 241 242 243 244
    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://'):
245
          title = self.save_url(from_nick, w)
246
          self.say_public(title)
Jonathan Harker's avatar
Jonathan Harker committed
247 248 249 250

  def say_public(self, text):
    "Print TEXT into public channel, for all to see."
    self.queue.send(text, self.channel)
251
    self.log_event(self.nickname, text)
Jonathan Harker's avatar
Jonathan Harker committed
252 253 254 255 256 257 258 259 260 261 262 263 264

  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
265
    "Return a random pondering."
Jonathan Harker's avatar
Jonathan Harker committed
266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289
    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()

    if cmd == 'help':
      self.reply(self.helptext, target)

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

    elif cmd == 'urls' or cmd == 'list':
290 291 292
      db = self.get_session()
      for url in db.query(Url).order_by(Url.timestamp):
        line = "%s %s" % (url.url, url.title)
293 294
        self.reply(line, target)
        time.sleep(1)
Jonathan Harker's avatar
Jonathan Harker committed
295 296

    elif cmd.startswith('http:') or cmd.startswith('https:'):
297
      title = self.save_url(from_private, cmd)
298 299 300 301
      if title == '':
        self.reply('URL added.', target)
      if title != '':
        self.reply('URL added: %s' % title, target)
Jonathan Harker's avatar
Jonathan Harker committed
302 303 304 305 306
      
    else:
      self.reply(self.exclaim(), target)


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 332 333 334 335 336 337
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
338 339
if __name__ == "__main__":
  try:
340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355
    (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
356 357 358
  except KeyboardInterrupt:
    print "Shutting down."