lolbot.py 9.87 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 81 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
    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))
  url       = Column(String(200))
  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 197 198 199 200 201 202 203 204 205 206 207 208
  def save_url(self, nickname, url):
    theurl = Url(nickname, url)
    db = self.get_session()
    db.add(theurl)
    db.commit()
    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
209

Jonathan Harker's avatar
Jonathan Harker committed
210 211 212 213 214 215 216 217 218 219 220 221 222 223 224
  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."

225
    # log it
Jonathan Harker's avatar
Jonathan Harker committed
226
    from_nick = nm_to_n(event.source())
227 228
    self.log_event(from_nick, event.arguments()[0])

Jonathan Harker's avatar
Jonathan Harker committed
229 230 231
    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)
232 233 234 235 236
    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://'):
237
          title = self.save_url(from_nick, w)
238
          self.say_public(title)
Jonathan Harker's avatar
Jonathan Harker committed
239 240 241 242

  def say_public(self, text):
    "Print TEXT into public channel, for all to see."
    self.queue.send(text, self.channel)
243
    self.log_event(self.nickname, text)
Jonathan Harker's avatar
Jonathan Harker committed
244 245 246 247 248 249 250 251 252 253 254 255 256

  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
257
    "Return a random pondering."
Jonathan Harker's avatar
Jonathan Harker committed
258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281
    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':
282 283 284
      db = self.get_session()
      for url in db.query(Url).order_by(Url.timestamp):
        line = "%s %s" % (url.url, url.title)
285 286
        self.reply(line, target)
        time.sleep(1)
Jonathan Harker's avatar
Jonathan Harker committed
287 288

    elif cmd.startswith('http:') or cmd.startswith('https:'):
289
      title = self.save_url(from_private, cmd)
290 291 292 293
      if title == '':
        self.reply('URL added.', target)
      if title != '':
        self.reply('URL added: %s' % title, target)
Jonathan Harker's avatar
Jonathan Harker committed
294 295 296 297 298
      
    else:
      self.reply(self.exclaim(), target)


299 300 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
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
330 331
if __name__ == "__main__":
  try:
332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347
    (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
348 349 350
  except KeyboardInterrupt:
    print "Shutting down."