From 06499d85169ce7322359bef83692d3c47416b10d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nikola=20Forr=C3=B3?= <nikola@mraveniste.cc>
Date: Fri, 20 Jul 2018 11:58:36 +0200
Subject: [PATCH] Make cooldowns per-server

---
 clients/discord.py | 26 ++++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

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: