rerunbot/rerunbot.py

import configparser
import datetime
import functools
import logging
import os
import re
import unicodedata

import dateutil.parser
import irc.bot
import requests
import tempora.schedule


class ReplayBot(irc.bot.SingleServerIRCBot):
    def __init__(self):
        self.logger = logging.getLogger('irc.client')
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter('%(levelname)s: %(message)s'))
        self.logger.addHandler(handler)
        level = logging.DEBUG if bool(int(os.getenv('DEBUG', 0))) else logging.INFO
        self.logger.setLevel(level)
        self.config = configparser.ConfigParser()
        self.config.read('settings.cfg')
        self.supported_commands = [
            (re.compile(r'^!load\s+(?P<id>\d+)$'), self._do_load),
            (re.compile(r'^!start(\s+(?P<at>.+))?$'), self._do_start),
            (re.compile(r'^!resync\s+(?P<offset>.+)$'), self._do_resync),
            (re.compile(r'^!stop$'), self._do_stop),
        ]
        self.messages = []
        self.output_channel = '#{0}'.format(self.config['IRC'].get('output_channel'))
        control_channel = self.config['IRC'].get('control_channel')
        self.control_channel = '#{0}'.format(control_channel) if control_channel else None
        self.control_users = self.config['IRC'].get('control_users').split(',')
        server = self.config['IRC'].get('server')
        port = self.config['IRC'].getint('port')
        nickname = self.config['IRC'].get('nickname')
        token = self.config['Twitch'].get('token')
        self.logger.info('Connecting to %s:%d', server, port)
        super(ReplayBot, self).__init__([(server, port, token)], nickname, nickname)

    def on_welcome(self, connection, event):
        connection.cap('REQ', ':twitch.tv/membership')
        connection.cap('REQ', ':twitch.tv/tags')
        connection.cap('REQ', ':twitch.tv/commands')
        channels = set([self.output_channel])
        if self.control_channel:
            channels.add(self.control_channel)
        for channel in channels:
            self.logger.info('Joining %s', channel)
            connection.join(channel)

    def on_join(self, connection, event):
        self.logger.info('Joined %s', event.target)

    def on_pubmsg(self, connection, event):
        self._process_message(connection, event)

    def on_whisper(self, connection, event):
        self._process_message(connection, event)

    def _send_response(self, connection, event, msg):
        if event.target.startswith('#'):
            connection.privmsg(event.target, msg)
        else:
            connection.privmsg('#jtv', '/w {0} {1}'.format(event.source.nick, msg))

    def _process_message(self, connection, event):
        tags = {t['key']: t['value'] for t in event.tags}
        message = ''.join([c for c in event.arguments[0] if not unicodedata.category(c).startswith('C')])
        message = message.rstrip()
        send_response = functools.partial(self._send_response, connection, event)
        for pattern, action in self.supported_commands:
            m = pattern.match(message)
            if m:
                action(tags, send_response, **m.groupdict())

    def _do_load(self, tags, send_response, id, **kwargs):
        if tags['user-id'] not in self.control_users:
            send_response('Sorry @{0}, you are not allowed to do this'.format(tags['display-name']))
            return
        client_id = self.config['Twitch'].get('client_id')
        user_blacklist = self.config['Rerun'].get('user_blacklist').split(',')
        def get_comments(cursor):
            url = 'https://api.twitch.tv/v5/videos/{0}/comments'.format(id)
            params = dict(client_id=client_id, cursor=cursor)
            r = requests.get(url, params=params)
            r.raise_for_status()
            return r.json()
        self.messages = []
        cursor = ''
        while True:
            data = get_comments(cursor)
            for comment in data.get('comments', []):
                if comment['commenter']['_id'] in user_blacklist:
                    continue
                self.messages.append((
                    datetime.timedelta(seconds=float(comment['content_offset_seconds'])),
                    comment['commenter']['display_name'],
                    comment['message']['body']))
            cursor = data.get('_next')
            if not cursor:
                break
        if self.messages:
            send_response('@{0}: loaded {1} messages, first message at {2}'.format(
                tags['display-name'], len(self.messages), str(self.messages[0][0])))
        else:
            send_response('@{0}: failed to load messages'.format(tags['display-name']))

    def _do_start(self, tags, send_response, at=None, **kwargs):
        if tags['user-id'] not in self.control_users:
            send_response('Sorry @{0}, you are not allowed to do this'.format(tags['display-name']))
            return
        if at is None:
            t = datetime.datetime.now().astimezone(tz=None)
        else:
            t = dateutil.parser.parse(at).astimezone(tz=None)
        for offset, user, msg in self.messages:
            def cb(user=user, msg=msg):
                return self._post_message(user, msg)
            self.reactor.scheduler.execute_at(t + offset, cb)
        with self.reactor.mutex:
            while self.reactor.scheduler.queue:
                if self.reactor.scheduler.queue[0].due():
                    self.reactor.scheduler.queue.pop(0)
                else:
                    break

    def _do_resync(self, tags, send_response, offset, **kwargs):
        if tags['user-id'] not in self.control_users:
            send_response('Sorry @{0}, you are not allowed to do this'.format(tags['display-name']))
            return
        offset = datetime.timedelta(seconds=float(offset))
        with self.reactor.mutex:
            for i, dc in enumerate(self.reactor.scheduler.queue):
                ndc = tempora.schedule.DelayedCommand.from_datetime(dc + offset)
                ndc.delay = dc.delay + offset
                ndc.target = dc.target
                self.reactor.scheduler.queue[i] = ndc
            while self.reactor.scheduler.queue:
                if self.reactor.scheduler.queue[0].due():
                    self.reactor.scheduler.queue.pop(0)
                else:
                    break

    def _do_stop(self, tags, send_response, **kwargs):
        if tags['user-id'] not in self.control_users:
            send_response('Sorry @{0}, you are not allowed to do this'.format(tags['display-name']))
            return
        with self.reactor.mutex:
            self.reactor.scheduler.queue.clear()

    def _post_message(self, user, msg):
        text = self.config['Rerun'].get('message_template').format(username=user, message=msg)
        self.connection.privmsg(self.output_channel, text)


def main():
    bot = ReplayBot()
    bot.start()


if __name__ == "__main__":
    main()