diff --git a/clients/discord.py b/clients/discord.py index e008f1f..cc74970 100644 --- a/clients/discord.py +++ b/clients/discord.py @@ -8,22 +8,23 @@ from commands import CommandError def cooldown(retries, timeout, failure): def do_cooldown(function): - def wrapper(self, user, *args, **kwargs): + def wrapper(self, server, user, *args, **kwargs): cooldowns = getattr(function, 'cooldowns', {}) - if user not in cooldowns: - cooldowns[user] = dict(tries=0, last_run=0) - cd = cooldowns[user] + 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, user, *args, **kwargs) + 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, user, *args, **kwargs) - return failure(self, user, *args, **kwargs) + return function(self, server, user, *args, **kwargs) + return failure(self, server, user, *args, **kwargs) return wrapper return do_cooldown @@ -52,15 +53,16 @@ class DiscordClient(discord.Client): for pattern, action in self.supported_commands: m = pattern.match(message.content) if m: - await action(message.author.id, message, **m.groupdict()) + server = message.server.id if message.server else None + await action(server, message.author.id, message, **m.groupdict()) - async def _cooldown_failure(self, user, message, **kwargs): + 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)) @cooldown(retries=2, timeout=5*60, failure=_cooldown_failure) - async def _do_bellagram(self, user, message, **kwargs): + async def _do_bellagram(self, server, user, message, **kwargs): try: bellagram = self.commands.bellagram() except CommandError as e: @@ -73,7 +75,7 @@ class DiscordClient(discord.Client): await self.send_message(message.channel, embed=embed) @cooldown(retries=3, timeout=5*60, failure=_cooldown_failure) - async def _do_yt(self, user, message, query, **kwargs): + async def _do_yt(self, server, user, message, query, **kwargs): try: result = self.commands.query_youtube(query) except CommandError as e: @@ -90,7 +92,7 @@ class DiscordClient(discord.Client): 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, user, message, filter, **kwargs): + async def _do_clip(self, server, user, message, filter, **kwargs): try: result = self.commands.find_clip(filter) except CommandError as e: