import datetime import re import time import discord from commands import CommandError def cooldown(retries, timeout, failure): def do_cooldown(function): def wrapper(self, server, user, *args, **kwargs): cooldowns = getattr(function, 'cooldowns', {}) hash = '{0}:{1}'.format(server, user) if hash not in cooldowns: cooldowns[hash] = dict(tries=0, last_run=0) cd = cooldowns[hash] if cd['tries'] < retries: cd['tries'] += 1 cd['last_run'] = time.time() setattr(function, 'cooldowns', cooldowns) return function(self, server, user, *args, **kwargs) if time.time() - cd['last_run'] > timeout: cd['tries'] = 1 cd['last_run'] = time.time() setattr(function, 'cooldowns', cooldowns) return function(self, server, user, *args, **kwargs) return failure(self, server, user, *args, **kwargs) return wrapper return do_cooldown class DiscordClient(discord.Client): def __init__(self, config, logger, commands, extra_commands): self.config = config self.logger = logger self.commands = commands self.extra_commands = extra_commands self.supported_commands = [ (re.compile(r'^!lastquote$'), self._do_lastquote), (re.compile(r'^!findquote\s+(?P")?(?P.+)(?(q)")$'), self._do_findquote), (re.compile(r'^!(bella(gram|pics)|insta(gram|bella))$'), self._do_bellagram), (re.compile(r'^!yt\s+(?P")?(?P.+)(?(q)")$'), self._do_yt), (re.compile(r'^!clip\s+(?P")?(?P.+)(?(q)")$'), self._do_clip), ] super(DiscordClient, self).__init__() async def start_(self): token = self.config['Discord'].get('token') await self.start(token) async def on_ready(self): self.logger.info('Logged in as {0}'.format(self.user.name)) async def on_message(self, message): await self._process_announcement(message) server = message.server.id if message.server else None for pattern, action in self.supported_commands: m = pattern.match(message.content) if m: await action(server, message.author.id, message, **m.groupdict()) for cmd, resp in self.extra_commands.items(): if cmd == message.content: await self.send_message(message.channel, resp.format(user=message.author.mention)) async def _process_announcement(self, message): announcer = self.config['Discord'].get('announcer') channels = self.config['Discord'].get('announcement_channels').split(',') if message.author.id != announcer or message.channel.id not in channels: return after = datetime.datetime.utcnow() - datetime.timedelta(minutes=5) for channel in [discord.Object(c) for c in channels if c != message.channel.id]: async for msg in self.logs_from(channel, after=after): if msg.content == message.content: return embed = next(iter(message.embeds or []), None) await self.send_message(channel, message.content, tts=message.tts, embed=embed) # TODO: twitter async def _cooldown_failure(self, server, user, message, **kwargs): await self.send_message(message.channel, 'Sorry {0}, you have to wait a while before running ' 'the same command again'.format(message.author.mention)) def _format_quote(self, quote): if '"' not in quote['text']: quote['text'] = '"{0}"'.format(quote['text']) return 'One time, Lilia said this... `#{id}: {text} [{game}] [{date}]`'.format(**quote) async def _do_lastquote(self, server, user, message, **kwargs): try: result = self.commands.last_quote() except CommandError as e: await self.send_message(message.channel, 'Sorry {0}, {1}'.format(message.author.mention, e)) else: await self.send_message(message.channel, self._format_quote(result)) async def _do_findquote(self, server, user, message, filter, **kwargs): try: result = self.commands.find_quote(filter) except CommandError as e: await self.send_message(message.channel, 'Sorry {0}, {1}'.format(message.author.mention, e)) else: await self.send_message(message.channel, self._format_quote(result)) @cooldown(retries=2, timeout=5*60, failure=_cooldown_failure) async def _do_bellagram(self, server, user, message, **kwargs): try: bellagram = self.commands.bellagram() except CommandError as e: await self.send_message(message.channel, 'Sorry {0}, {1}'.format(message.author.mention, e)) else: embed = discord.Embed(title=bellagram['title'], url=bellagram['url'], color=0xd2a517) embed.set_image(url=bellagram['display_url']) embed.set_author(name=bellagram['owner'], url=bellagram['owner_url'], icon_url=bellagram['owner_pic_url']) await self.send_message(message.channel, embed=embed) @cooldown(retries=3, timeout=5*60, failure=_cooldown_failure) async def _do_yt(self, server, user, message, query, **kwargs): try: result = self.commands.query_youtube(query) except CommandError as e: await self.send_message(message.channel, 'Sorry {0}, {1}'.format(message.author.mention, e)) else: if result['kind'] == 'playlist': embed = discord.Embed(title=result['title'], url=result['url'], description=result['description'], color=0xff0000) embed.set_thumbnail(url=result['thumbnail_url']) embed.set_author(name=result['channel_title'], url=result['channel_url'], icon_url=result['channel_thumbnail_url']) await self.send_message(message.channel, embed=embed) else: await self.send_message(message.channel, '{title}\n{url}'.format(**result)) @cooldown(retries=3, timeout=5*60, failure=_cooldown_failure) async def _do_clip(self, server, user, message, filter, **kwargs): try: result = self.commands.find_clip(filter) except CommandError as e: await self.send_message(message.channel, 'Sorry {0}, {1}'.format(message.author.mention, e)) else: await self.send_message(message.channel, '{title}\n{url}'.format(**result))