lolbot.py 16.7 KB
Newer Older
1
#! /usr/bin/env python
Jonathan Harker's avatar
Jonathan Harker committed
2
"""
3 4 5 6 7
LOLBOT 2

 - die: Let the bot cease to exist.
 - ask: Ask a MoxQuizz question.
 - list: list some URLs
Jonathan Harker's avatar
Jonathan Harker committed
8 9
"""

10
from __future__ import print_function, unicode_literals
11

12
import sys
13
import os
14 15 16
import sqlite3
import random
import time
17
import getopt
18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35
import irc.strings
from irc.bot import SingleServerIRCBot
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from models import Log, Url, Model
from pymoxquizz import QuestionBank, Question
from os import listdir, path


DEBUG = True


def debug(msg):
    if DEBUG:
        print(msg)


class LolBot(SingleServerIRCBot):
36
    """
37 38
    An IRC bot to entertain the troops with MoxQuizz questions, log URLs, and
    other shenanigans.
39 40
    """

41 42
    qb = list()

43 44 45 46 47 48 49 50 51 52 53 54
    def __init__(self, config=None):
        """
        Constructor. Instantiates a lolbot with a configuration dictionary,
        or from command-line options if none is specified.
        """

        if not config:
            config = LolBot.get_options()

        if not self.validate_config(config):
            sys.exit(1)

55
        self.config = config
56 57 58 59 60
        (server, port, channel, nickname, database) = (
            config['irc.server'],
            config['irc.port'],
            config['irc.channel'],
            config['irc.nickname'],
61
            config['db.file'])
62

63
        debug("Instantiating SingleServerIRCBot")
64
        irc.client.ServerConnection.buffer_class = irc.buffer.LenientDecodingLineBuffer
65 66
        SingleServerIRCBot.__init__(self, [(server, port)], nickname, nickname)
        self.channel = channel
Brett Wilkins's avatar
Brett Wilkins committed
67

68
        # load some MoxQuizz questions
69 70 71 72 73 74
        qfiles = [f for f in listdir('questions') if path.isfile(path.join('questions', f))]
        debug("Loading MoxQuizz questions")
        for f in qfiles:
            qfile = path.abspath(path.join('questions', f))
            debug("  - from MoxQuizz bank '%s'" % qfile)
            self.qb += QuestionBank(qfile).questions
75
        random.shuffle(self.qb)
76
        self.quiz = 0
77 78
        self.question = None

79
        # connect to the database
80 81
        debug("Connecting to SQLite database '%s'" % database)
        self.dbfile = database
Brett Wilkins's avatar
Brett Wilkins committed
82
        self.dbengine = create_engine('sqlite+pysqlite://', creator=self._get_connection)
83 84
        Model.metadata.bind = self.dbengine
        Model.metadata.create_all()
85
        self.get_db = sessionmaker(bind=self.dbengine)
86

87 88
        self.helptext = "Keeps a list of URLs. Commands: list [n|x-y] - prints the last 10 URLs (or n URLs, or x through y); <url> - adds the URL to the list; help - this message."
        debug("Exiting lolbot constructor")
89

90 91 92 93 94 95
    def _get_connection(self):
        """Creator function for SQLAlchemy."""
        connection = sqlite3.Connection(self.dbfile)
        connection.text_factory = str
        debug("Creating SQLAlchemy connection")
        return connection
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
    @property
    def nickname(self):
        return self.connection.get_nickname()

    def say_public(self, text):
        """
        Say text in the public channel for all to see.
        """
        self.connection.privmsg(self.channel, text)
        self.log_event(self.nickname, text)

    def say_private(self, nick, text):
        """
        Say text in a private message to nick.
        """
        self.connection.privmsg(nick, text)

    def reply(self, text, to_private=None):
        """
        Say text in either public channel or a private message (if to_private
        supplied).
        """
        if to_private is not None:
            self.say_private(to_private, text)
        else:
            self.say_public(text)

124
    def on_nicknameinuse(self, connection, event):
125 126
        debug("Nick '%s' in use, trying '%s_'" % (self.nickname, self.nickname))
        connection.nick(self.nickname + "_")
127

128
    def on_welcome(self, connection, event):
129
        debug("Joining channel '%s' as %s" % (self.channel, self.nickname))
130
        connection.join(self.channel)
131

132
    def on_privmsg(self, connection, event):
133 134 135 136 137 138
        """
        Handle a /msg from a user.
        Handle commands addressed to the bot.
        """
        message = event.arguments[0]
        self.do_command(event, message)
139

140
    def on_pubmsg(self, connection, event):
141
        """
142 143 144
        Handle an event on the channel.
        Handle commands addressed to the bot.
        If there's a question, see if it's been answered.
145
        """
146 147

        # Handle bot commands if addressed by nick or using ! shortcut.
148
        try:
149
            (nick, message) = event.arguments[0].split(": ", 1)
150 151 152 153 154
            if irc.strings.lower(nick) == irc.strings.lower(self.connection.get_nickname()):
                self.do_command(event, message.strip())
        except ValueError:
            message = event.arguments[0]
            nick = event.source.nick
155 156
            if message.startswith('!'):
                self.do_command(event, message.lstrip('!'))
157

158 159 160 161
        # Log it.
        self.log_event(nick, message)

        # Deal with MoxQuizz question.
162 163 164
        if self.quiz:
            self.handle_quiz(nick, message)

165 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 192 193 194 195 196 197 198 199 200
        # Record URLs.
        words = message.split(" ")
        for w in words:
            if w.startswith('http://') or w.startswith('https://'):
                title = self.save_url(nick, w)
                if title is False:
                    self.say_public("Failed to record URL, or no title found.")
                else:
                    self.say_public("URL added. %s" % title)

    def save_url(self, nickname, url):
        title = False
        try:
            db = self.get_db()
            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 as ex:
            print("Exception caught saving URL: %s" % ex)
        return title

    def log_event(self, nick, text):
        try:
            entry = Log(nick, text)
            db = self.get_db()
            db.add(entry)
            db.commit()
            print(entry)
        except Exception as ex:
            print("Exception caught logging event: %s" % ex)

201 202 203 204 205 206 207 208 209 210
    def start_quiz(self, nick):
        self.quiz = 0
        self.quiz_scores = dict()
        self.connection.notice(self.channel, 'Quiz begun by %s.' % nick)
        self.quiz_get_next()

    def stop_quiz(self):
        self.quiz = 0
        self.quiz_scores = None
        self.question = None
211

212 213 214 215
    def quiz_get_next(self):
        self.quiz += 1
        self.question = random.choice(self.qb)
        print(str(self.question.question))
216
        self.connection.notice(self.channel, "Question %s: %s" % (self.quiz, str(self.question.question)))
217 218 219 220 221 222

    def quiz_award_points(self, nick):
        if nick not in self.quiz_scores.keys():
            self.quiz_scores[nick] = 0
        self.quiz_scores[nick] += self.question.score

223 224 225 226 227 228 229 230
        score = "%s point" % self.quiz_scores[nick]
        if self.quiz_scores[nick] != 1:
            score += "s"

        self.connection.notice(self.channel, 'Correct! The answer was %s.' % self.question.answer)
        time.sleep(1)
        self.connection.notice(self.channel, '%s is on %s.' % (nick, score))

231 232 233 234 235 236 237
    def quiz_check_win(self, nick):
        if self.quiz_scores[nick] == 10:
            self.connection.notice(self.channel, '%s wins with 10 points!' % nick)
            self.quiz_scoreboard()
            self.stop_quiz()

    def quiz_scoreboard(self):
238 239 240 241
        if not self.quiz:
            self.connection.notice(self.channel, 'Quiz not running.')
            return

242 243 244 245 246 247
        self.connection.notice(self.channel, 'Scoreboard:')
        for nick in self.quiz_scores.keys():
            score = "%s point" % self.quiz_scores[nick]
            if self.quiz_scores[nick] != 1:
                score += "s"
            self.connection.notice(self.channel, '%s has %s.' % (nick, score))
248 249
        if not len(self.quiz_scores):
            self.connection.notice(self.channel, 'So far, nobody has got anything right.')
250 251 252 253 254 255 256 257 258

    def handle_quiz(self, nick, message):
        # bail if there's no quiz or unanswered question.
        if not self.quiz or not isinstance(self.question, Question):
            return

        # see if anyone answered correctly.
        if self.question.attempt(message):
            self.quiz_award_points(nick)
259
            time.sleep(1)
260
            self.quiz_check_win(nick)
261

262 263
            # if nobody has won, carry on
            if self.quiz:
264 265 266 267 268

                # scores every 10 questions.
                if self.quiz % 10 == 0:
                    self.quiz_scoreboard()

269
                time.sleep(1)
270 271 272 273 274 275 276 277
                self.quiz_get_next()

    def do_command(self, e, cmd):
        """
        Handle bot commands.
        """
        nick = e.source.nick
        c = self.connection
278

279 280
        if cmd == "die":
            self.die()
281

282 283
        elif cmd == 'help':
            c.notice(nick, self.helptext)
284

285
        elif cmd == 'status':
286 287 288 289
            c.notice(self.channel, "I know %s questions." % len(self.qb))
            if self.quiz:
                c.notice(self.channel, "I am currently running a quiz.")
                self.quiz_scoreboard()
290

291 292 293 294 295
        elif cmd == 'halt' or cmd == 'quit':
            if self.quiz:
                self.quiz_scoreboard()
                self.stop_quiz()
                c.notice(self.channel, "Quiz halted by %s. Use ask to start a new one." % nick)
296
            else:
297 298
                c.notice(self.channel, "No quiz running.")

299 300 301 302
        elif cmd == 'scores':
            if self.quiz:
                self.quiz_scoreboard()

303 304 305 306 307 308 309 310 311
        elif cmd == 'ask':
            if self.quiz:
                c.notice(self.channel, "Quiz is running. Use halt or quit to stop.")
                c.notice(self.channel, str(self.question.question))
            elif isinstance(self.question, Question):
                c.notice(self.channel, "There is an unanswered question.")
                c.notice(self.channel, str(self.question.question))
            else:
                self.start_quiz(nick)
312

313 314 315 316
        elif cmd == 'revolt':
            if isinstance(self.question, Question):
                c.notice(self.channel, "Fine, the answer is: %s" % self.question.answer)
                self.quiz_get_next()
317

318
        elif cmd.startswith('urls') or cmd.startswith('list'):
319
            db = self.get_db()
320
            try:
321
                (listcmd, n) = cmd.split(" ", 1)
322
            except ValueError:
323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350
                n = '5'

            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 as ex:
                    c.notice(nick, "Give me a number or a range of numbers, e.g. list 5 or list 11-20")
                    raise ex
                rows = db.query(Url).order_by(Url.timestamp.desc())[x - 1: y]
            else:
                try:
                    n = abs(int(n))
                except ValueError as ex:
                    c.notice(nick, "Give me a number or a range of numbers, e.g. list 5 or list 11-20")
                    raise ex
                rows = db.query(Url).order_by(Url.timestamp.desc())[:n]

            for url in rows:
                line = "%s %s" % (url.url, url.title)
                c.notice(nick, line)
                time.sleep(1)
351

352 353
        else:
            c.notice(nick, "Not understood: " + cmd)
354

355 356 357 358 359 360 361 362 363 364
    def validate_config(self, config):
        """
        Basic checks for configuration parameters. Returns a Boolean indicating
        success or failure.
        """

        # validate IRC host
        if 'irc.server' not in config.keys():
            print("Error: the IRC server was not specified. Use --help for more information.")
            return False
365

366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394
        # validate IRC port
        if 'irc.port' not in config.keys():
            config['irc.port'] = '6667'
            try:
                config['irc.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.")
                return False

        # 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.")
            return False

        # validate bot nickname
        if 'irc.nickname' not in config.keys():
            config['irc.nickname'] = 'lolbot'

        # validate bot nickname
        if 'db.file' not in config.keys():
            config['db.file'] = 'lolbot.db'

        return True

    @staticmethod
    def get_options():
        """
        Set up configuration from the script arguments.
        """
395

396
        try:
397 398 399 400 401 402 403 404 405 406 407 408 409 410 411
            (options, args) = getopt.getopt(sys.argv[1:], 'hc:s:p:j:n:d:', ['help', 'config=', 'server=', 'port=', 'join=', 'nick=', 'database=', ])
        except getopt.GetoptError as err:
            print(str(err))
            LolBot.usage()
            sys.exit(2)

        config = {}
        for option, value in options:
            # Display help text.
            if option in ('-h', '--help'):
                LolBot.usage()
                sys.exit(2)

            # Get configuration from a file.
            if option in ('-c', '--config'):
412
                config = LolBot.load_config(value)
413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446
                break

            # Individually specified settings.
            if option in ('-s', '--server'):
                config['irc.server'] = value
            if option in ('-p', '--port'):
                config['irc.port'] = value
            if option in ('-j', '--join', '--channel', '--join-channel'):
                config['irc.channel'] = value
            if option in ('-n', '--nickname'):
                config['irc.nickname'] = value
            if option in ('-d', '--database'):
                config['db.file'] = value

        return config

    @staticmethod
    def load_config(config_path=''):
        """
        This method loads configuration options from a lolbot.conf file. The file
        should look something like this::

        irc.server = irc.yourdomain.com
        irc.port = 6667
        irc.channel = #lolbottest
        irc.nickname = lolbot
        db.file = lolbot.db
        """

        if 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.")
            LolBot.usage()
447
            sys.exit(1)
448

449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500
        # 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

        return config

    @staticmethod
    def usage():
        """
        Spits out CLI help.
        """

        print("""Run a lolbot.

      -h, --help
             This message.

      -s, --server=<hostname>
             The IRC server, e.g. irc.freenode.net

      -p, --port=<port>
             The IRC server port. Default: 6667

      -j, --join=<channel>
             The chat channel to join, e.g. #trainspotting

      -n, --nickname=<nickname>
             The nickname for your lolbot. Default: "lolbot"

      -d, --database=<path>
             The path to a SQLite lolbot database. If not specified, lolbot will
             attempt to use a SQLite database called lolbot.db in the same
             directory as the script.

    Configuration file:

      -c, --config=<filename>
             Specify a configuration file. Ignores any options specified from the
             command-line. Defaults to lolbot.conf in the same directory as the
             script. File layout:
501

502 503 504 505 506 507
               irc.server = <host>
               irc.port = <port>
               irc.channel = <channel>
               irc.nickname = <nickname>
               db.file = <path>
""".strip())
508

509

Jonathan Harker's avatar
Jonathan Harker committed
510
if __name__ == "__main__":
511
    try:
512
        LolBot().start()
513 514
    except KeyboardInterrupt:
        pass