From cd28334e789726aa0b2538e1ca64f2b4cfb50edb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nikola=20Forr=C3=B3?= <nikola@mraveniste.cc>
Date: Tue, 22 May 2018 18:10:32 +0200
Subject: [PATCH] Implement per-user cooldown on Discord

---
 clients/discord.py   | 20 +++++++++++++++++++-
 settings.cfg.example |  3 +++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/clients/discord.py b/clients/discord.py
index a4e98bf..d88b5ff 100644
--- a/clients/discord.py
+++ b/clients/discord.py
@@ -1,4 +1,5 @@
 import re
+import time
 
 import discord
 
@@ -15,6 +16,7 @@ class DiscordClient(discord.Client):
             (re.compile(r'^!(bella(gram|pics)|insta(gram|bella))$'), self._do_bellagram),
             (re.compile(r'^!yt\s+(?P<q>")?(?P<query>.+)(?(q)")$'), self._do_yt),
         ]
+        self.cooldowns = {p.pattern: {} for p, _ in self.supported_commands}
         super(DiscordClient, self).__init__()
 
     async def start_(self):
@@ -25,10 +27,26 @@ class DiscordClient(discord.Client):
         self.logger.info('Logged in as {0}'.format(self.user.name))
 
     async def on_message(self, message):
+        def check_cooldown(pattern, user):
+            retries = self.config['Discord'].getint('cooldown_retries')
+            seconds = self.config['Discord'].getint('cooldown_seconds')
+            if user not in self.cooldowns[pattern]:
+                self.cooldowns[pattern][user] = dict(tries=0, time=0)
+            cooldown = self.cooldowns[pattern][user]
+            if cooldown['tries'] >= retries and time.time() - cooldown['time'] <= seconds:
+                return False
+            cooldown['tries'] = 1 if cooldown['tries'] >= retries else cooldown['tries'] + 1
+            cooldown['time'] = time.time()
+            return True
         for pattern, action in self.supported_commands:
             m = pattern.match(message.content)
             if m:
-                await action(message, **m.groupdict())
+                if check_cooldown(pattern.pattern, message.author.id):
+                    await action(message, **m.groupdict())
+                else:
+                    await self.send_message(message.channel,
+                        'Sorry {0}, you have to wait a while before running '
+                        'the same command again'.format(message.author.mention))
 
     async def _do_bellagram(self, message, **kwargs):
         try:
diff --git a/settings.cfg.example b/settings.cfg.example
index d7e7614..f2c9db1 100644
--- a/settings.cfg.example
+++ b/settings.cfg.example
@@ -6,6 +6,9 @@ channels = lilialil
 
 [Discord]
 token = __DISCORD_TOKEN__
+# 5 minutes cooldown after 2 retries
+cooldown_retries = 2
+cooldown_seconds = 300
 
 [Twitch]
 api_url = https://api.twitch.tv/v5