lolbot.py 12.5 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 13
from __future__ import print_function  # unicode_literals

14
try:
15
    import sys
16
    import os
17 18 19 20
    import string
    import random
    import time
    import getopt
Brett Wilkins's avatar
Brett Wilkins committed
21
    import sqlite3
22 23 24
    from twisted.words.protocols import irc
    from twisted.internet import protocol
    from twisted.internet import reactor
25
    from datetime import datetime
26
    from sqlalchemy import create_engine
27
    from sqlalchemy.orm import sessionmaker
28 29
    from pymoxquizz import QuestionBank, Question
    from models import Url, Log, Model
30
except ImportError:
31
    print("Some modules missing: Lolbot relies on Twisted IRC, Mechanize and SQLAlchemy.\n")
32
    sys.exit
33

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

52 53 54 55 56 57

class LolBot(irc.IRCClient):
    """
    The Lolbot itself.
    """

Brett Wilkins's avatar
Brett Wilkins committed
58 59 60 61 62
    def _get_connection(self):
        connection = sqlite3.Connection(self.config['db.file'])
        connection.text_factory = str
        return connection

63
    def created(self, when):
64 65 66 67 68
        # load some MoxQuizz questions
        self.qb = QuestionBank('/home/johnno/questions.doctorlard.en').questions
        random.shuffle(self.qb)
        self.question = None

69
        # connect to the database
Brett Wilkins's avatar
Brett Wilkins committed
70
        self.dbengine = create_engine('sqlite+pysqlite://', creator=self._get_connection)
71 72
        Model.metadata.bind = self.dbengine
        Model.metadata.create_all()
73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89
        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()
90
            print(theurl)
91
            title = theurl.title
92 93
        except Exception as ex:
            print("Exception caught saving URL: %s" % ex)
94 95 96 97 98 99 100 101
        return title

    def log_event(self, nick, text):
        try:
            entry = Log(nick, text)
            db = self.get_session()
            db.add(entry)
            db.commit()
102 103 104
            print(entry)
        except Exception as ex:
            print("Exception caught logging event: %s" % ex)
105 106 107 108 109 110 111 112 113 114 115 116 117 118 119

    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)
120
        print("Signed on as %s." % (self.nickname,))
121 122

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

    def privmsg(self, user, channel, msg):
Jonathan Harker's avatar
Jonathan Harker committed
126
        user = user.split('!')[0]
127

Jonathan Harker's avatar
Jonathan Harker committed
128
        if channel == self.nickname:
129 130 131 132 133 134
            # Private /msg from a user
            self.do_command(msg, user)
        else:
            # log it
            self.log_event(user, msg)

135 136 137 138 139 140 141 142 143
            # deal with MoxQuizz question
            answered = False
            if isinstance(self.question, Question) and self.question.answer.lower() in msg.lower():
                answered = True
                print("%s +1" % user)
                self.reply('Correct! %s scores one point.' % user)
            if answered:
                self.question = None

144 145 146 147 148 149 150 151 152
            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)
153
                        if title is False:
154 155 156 157 158 159
                            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
160
        self.notice(self.channel, text)
161 162 163 164
        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
165
        self.notice(nick, text)
166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191

    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)

192 193 194 195
            elif cmd == 'ask':
                self.question = random.choice(self.qb)
                self.reply(str(self.question.question))

196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218
            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
219
                    except ValueError as ex:
220 221 222 223 224 225
                        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))
226
                    except ValueError as ex:
227 228 229 230 231 232 233 234 235 236
                        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:'):
237 238
                title = self.save_url(target, cmd)
                if title is False:
239 240 241 242 243 244 245
                    self.say_public("Sorry, I'm useless at UTF-8.")
                else:
                    self.reply('URL added. %s' % title, target)

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

246 247 248
        except Exception as ex:
            print("Exception caught processing command: %s" % ex)
            print("   command was '%s' from %s" % (cmd, target))
249 250 251
            self.reply("Sorry, I didn't understand:  %s" % cmd, target)
            self.reply(self.helptext, target)

252

253 254 255 256 257 258 259 260 261 262 263
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):
264
        print("Lost connection (%s), reconnecting." % (reason,))
265 266 267
        connector.connect()

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

270 271 272 273

def get_options():
    try:
        (options, args) = getopt.getopt(sys.argv[1:], "hc:", ["help", "config=", ])
274 275
    except getopt.GetoptError as err:
        print(str(err))
276 277 278 279
        usage()
        sys.exit(2)

    config_path = ""
280
    for option, value in options:
281 282 283 284 285
        if option in ("-h", "--help"):
            usage()
            sys.exit(2)
        if option in ("-c", "--config"):
            config_path = value
286 287
    return {'config_path': config_path}

288 289

def get_config(config_path):
290 291 292 293
    """
    This method loads configuration options from a lolbot.conf file. The file
    should look like this:

294 295 296 297 298
    irc.server = irc.freenode.net
    irc.port = 6667
    irc.channel = #lolbottest
    irc.nickname = lolbot
    db.url = sqlite:///lolbot.db
299 300 301 302

    """

    if not config_path:
303
        config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "lolbot.conf")
304
    if not os.path.exists(config_path):
305
        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.")
306 307
        usage()
        sys.exit(1)
308 309 310 311

    # open the configuration file and grab all param=value declarations.
    config = {}
    with open(config_path) as f:
312 313 314 315 316 317 318 319 320 321 322 323
        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
324 325 326

    # validate IRC host
    if "irc.server" not in config.keys():
327
        print("Error: the IRC server was not specified. Use --help for more information.")
328
        sys.exit(1)
329 330 331

    # validate IRC port
    if "irc.port" not in config.keys():
332
        config["irc.port"] = "6667"
333
    try:
334
        config["irc.port"] = int(config["irc.port"])
335
    except ValueError:
336
        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.")
337
        sys.exit(1)
338 339 340

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

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

348
    return config
Jonathan Harker's avatar
Jonathan Harker committed
349

350

351
def usage():
352
    print("""Run a lolbot.
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

  -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.
380 381
""")

382

Jonathan Harker's avatar
Jonathan Harker committed
383
if __name__ == "__main__":
384 385 386 387 388 389
    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:
390
        print("Shutting down.")