lolbot.py 13.6 KB
Newer Older
Jonathan Harker's avatar
Jonathan Harker committed
1 2 3 4
#!/usr/bin/env python
#
# LolBot
#
5
# New version based on Twisted IRC.
Jonathan Harker's avatar
Jonathan Harker committed
6 7 8

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

12
try:
13
    import sys
14
    import os
15 16 17 18
    import string
    import random
    import time
    import getopt
Brett Wilkins's avatar
Brett Wilkins committed
19
    import sqlite3
20 21 22
    from twisted.words.protocols import irc
    from twisted.internet import protocol
    from twisted.internet import reactor
23 24
    from datetime import datetime
    from mechanize import Browser
Brett Wilkins's avatar
Brett Wilkins committed
25
    from sqlalchemy import MetaData, Table, Column, String, Text, Integer, DateTime, create_engine
26 27 28
    from sqlalchemy.orm import sessionmaker
    from sqlalchemy.ext.declarative import declarative_base
except ImportError:
29
    print "Some modules missing: Lolbot relies on Twisted IRC, Mechanize and SQLAlchemy.\n"
30
    sys.exit
31

Jonathan Harker's avatar
Jonathan Harker committed
32 33 34 35 36 37 38 39 40 41 42 43 44 45
# 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
46 47
    "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
48 49
]

50 51 52
SqlBase = declarative_base()

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

58
    __tablename__ = "log"
59

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

65 66 67 68 69 70 71 72 73
    def __init__(self, nickname, text, timestamp=None):
        """
        Creates an event log for the IRC logger.
        """
        if timestamp is None:
            timestamp = datetime.now()
        self.timestamp = timestamp
        self.nickname = nickname
        self.text = text
74

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

class Url(SqlBase):
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
    """
    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), unique=True)
    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(irc.IRCClient):
    """
    The Lolbot itself.
    """

Brett Wilkins's avatar
Brett Wilkins committed
120 121 122 123 124
    def _get_connection(self):
        connection = sqlite3.Connection(self.config['db.file'])
        connection.text_factory = str
        return connection

125 126
    def created(self, when):
        # connect to the database
Brett Wilkins's avatar
Brett Wilkins committed
127
        self.dbengine = create_engine('sqlite+pysqlite://', creator=self._get_connection)
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 173 174 175 176 177 178 179 180 181 182
        SqlBase.metadata.bind = self.dbengine
        SqlBase.metadata.create_all()
        self.get_session = sessionmaker(bind=self.dbengine)

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

    def now(self):
        return datetime.today().strftime("%Y-%m-%d %H:%M:%S")

    def save_url(self, nickname, url):
        title = False
        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:
                theurl = db.query(Url).filter(Url.url == url).one()
            print theurl
            title = theurl.title
        except Exception, ex:
            print "Exception caught saving URL: %s" % ex
        return title

    def log_event(self, nick, text):
        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

    def _get_nickname(self):
        return self.factory.nickname
    nickname = property(_get_nickname)

    def _get_channel(self):
        return self.factory.channel
    channel = property(_get_channel)

    def _get_config(self):
        return self.factory.config
    config = property(_get_config)

    def signedOn(self):
        self.join(self.channel)
        print "Signed on as %s." % (self.nickname,)

    def joined(self, channel):
        print "Joined %s." % (channel,)

    def privmsg(self, user, channel, msg):
Jonathan Harker's avatar
Jonathan Harker committed
183 184
        user = user.split('!')[0]
        if channel == self.nickname:
185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206
            # Private /msg from a user
            self.do_command(msg, user)
        else:
            # log it
            self.log_event(user, msg)

            args = string.split(msg, ":", 1)
            if len(args) > 1 and args[0] == self.nickname:
                self.do_command(string.strip(args[1]))
            else:
                # parse it for links, add URLs to the list
                words = msg.split(" ")
                for w in words:
                    if w.startswith('http://') or w.startswith('https://'):
                        title = self.save_url(user, w)
                        if title == False:
                            self.say_public("Sorry, I'm useless at UTF-8.")
                        else:
                            self.say_public("URL added. %s" % title)

    def say_public(self, text):
        "Print TEXT into public channel, for all to see."
Jonathan Harker's avatar
Jonathan Harker committed
207
        self.notice(self.channel, text)
208 209 210 211
        self.log_event(self.nickname, text)

    def say_private(self, nick, text):
        "Send private message of TEXT to NICK."
Jonathan Harker's avatar
Jonathan Harker committed
212
        self.notice(nick, text)
213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 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 320 321 322 323 324 325 326 327 328 329

    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):
        "Return a random pondering."
        return random.choice(ponderings)

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

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

        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.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:
                    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 == False:
                    self.say_public("Sorry, I'm useless at UTF-8.")
                else:
                    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)
            self.reply("Sorry, I didn't understand:  %s" % cmd, target)
            self.reply(self.helptext, target)

class LolBotFactory(protocol.ClientFactory):
    protocol = LolBot

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

    def clientConnectionLost(self, connector, reason):
        print "Lost connection (%s), reconnecting." % (reason,)
        connector.connect()

    def clientConnectionFailed(self, connector, reason):
        print "Could not connect: %s" % (reason,)

def get_options():
    try:
        (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
    return { 'config_path': config_path }

def get_config(config_path):
330 331 332 333
    """
    This method loads configuration options from a lolbot.conf file. The file
    should look like this:

334 335 336 337 338
    irc.server = irc.freenode.net
    irc.port = 6667
    irc.channel = #lolbottest
    irc.nickname = lolbot
    db.url = sqlite:///lolbot.db
339 340 341 342

    """

    if not config_path:
343
        config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "lolbot.conf")
344
    if not os.path.exists(config_path):
345 346 347
        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)
348 349 350 351

    # open the configuration file and grab all param=value declarations.
    config = {}
    with open(config_path) as f:
352 353 354 355 356 357 358 359 360 361 362 363
        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
364 365 366

    # validate IRC host
    if "irc.server" not in config.keys():
367 368
        print "Error: the IRC server was not specified. Use --help for more information."
        sys.exit(1)
369 370 371

    # validate IRC port
    if "irc.port" not in config.keys():
372
        config["irc.port"] = "6667"
373
    try:
374
        config["irc.port"] = int(config["irc.port"])
375
    except ValueError:
376 377
        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)
378 379 380

    # validate IRC channel
    if "irc.channel" not in config.keys() or not config["irc.channel"].startswith("#"):
381 382
        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)
383 384 385

    # validate bot nickname
    if "irc.nickname" not in config.keys():
386
        config["irc.nickname"] = "lolbot"
Jonathan Harker's avatar
Jonathan Harker committed
387

388
    return config
Jonathan Harker's avatar
Jonathan Harker committed
389

390
def usage():
391
    print """Run a lolbot.
392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420

  -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
421
if __name__ == "__main__":
422 423 424 425 426 427 428
    args = get_options()
    config = get_config(args['config_path'])
    try:
        reactor.connectTCP(config['irc.server'], config['irc.port'], LolBotFactory(args['config_path']))
        reactor.run()
    except KeyboardInterrupt:
        print "Shutting down."
Jonathan Harker's avatar
Jonathan Harker committed
429