import re import time import discord from commands import CommandError def cooldown(retries, timeout, failure): def do_cooldown(function): def wrapper(self, user, *args, **kwargs): cooldowns = getattr(function, 'cooldowns', {}) if user not in cooldowns: cooldowns[user] = dict(tries=0, last_run=0) cd = cooldowns[user] if cd['tries'] < retries: cd['tries'] += 1 cd['last_run'] = time.time() setattr(function, 'cooldowns', cooldowns) return function(self, 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, user, *args, **kwargs) return failure(self, user, *args, **kwargs) return wrapper return do_cooldown class DiscordClient(discord.Client): def __init__(self, config, logger, commands): self.config = config self.logger = logger self.commands = commands self.supported_commands = [ (re.compile(r'^(?P!|\?)(bella(gram|pics)|insta(gram|bella))$'), self._do_bellagram), (re.compile(r'^(?P!|\?)yt\s+(?P")?(?P.+)(?(q)")$'), self._do_yt), ] 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): for pattern, action in self.supported_commands: m = pattern.match(message.content) if m: await action(message.author.id, message, **m.groupdict()) async def _cooldown_failure(self, 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)) @cooldown(retries=2, timeout=5*60, failure=_cooldown_failure) async def _do_bellagram(self, user, message, prefix, **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: if prefix == '!': embed = discord.Embed(title=bellagram['title'], url=bellagram['url'], color=0x8545bc) 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) elif prefix == '?': await self.send_message(message.channel, bellagram['url']) @cooldown(retries=3, timeout=5*60, failure=_cooldown_failure) async def _do_yt(self, user, message, query, prefix, **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 prefix == '!': 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) elif prefix == '?': await self.send_message(message.channel, result['url'])