Add global and channel rate limits

sopel/bot.py

@@ -420,31 +420,60 @@ def reply(self, message, destination=None, reply_to=None, notice=False):
 
     def call(self, func, sopel, trigger):
         nick = trigger.nick
+        current_time = time.time()
         if nick not in self._times:
             self._times[nick] = dict()
-
-        if not trigger.admin and \
-                not func.unblockable and \
-                func.rate > 0 and \
-                func in self._times[nick]:
-            timediff = time.time() - self._times[nick][func]
-            if timediff < func.rate:
-                self._times[nick][func] = time.time()
+        if self.nick not in self._times:
+            self._times[self.nick] = dict()
+
+        if not trigger.admin and not func.unblockable:
+            usertimediff = current_time - self._times[nick][func]
+            if func.user_rate > 0 and usertimediff < func.user_rate and \
+                    func in self._times[nick]:
+                self._times[nick][func] = current_time
+                LOGGER.info(
+                    "%s prevented from using %s in %s due to user limit: %d < %d",
+                    trigger.nick, func.__name__, trigger.sender, usertimediff,
+                    func.user_rate
+                )
+                return
+            globaltimediff = current_time - self._times[self.nick][func]
+            if func.global_rate > 0 and globaltimediff < func.global_rate and \
+                    func in self._times[self.nick]:
+                self._times[self.nick][func] = current_time
                 LOGGER.info(
-                    "%s prevented from using %s in %s: %d < %d",
-                    trigger.nick, func.__name__, trigger.sender, timediff,
-                    func.rate
+                    "%s prevented from using %s in %s due to global limit: %d < %d",
+                    trigger.nick, func.__name__, trigger.sender, globaltimediff,
+                    func.global_rate
                 )
                 return
 
+            if not trigger.is_privmsg:
+                chan = trigger.sender
+                if chan not in self._times:
+                    self._times[chan] = dict()
+                chantimediff = current_time - self._times[chan][func]
+                if func.channel_rate > 0 and chantimediff < func.channel_rate and \
+                            func in self._times[chan]:
+                        self._times[chan][func] = current_time
+                        LOGGER.info(
+                            "%s prevented from using %s in %s due to channel limit: %d < %d",
+                            trigger.nick, func.__name__, trigger.sender, timediff,
+                            func.channel_rate
+                        )
+                        return
+
         try:
             exit_code = func(sopel, trigger)
         except Exception:
             exit_code = None
             self.error(trigger)
 
         if exit_code != NOLIMIT:
-            self._times[nick][func] = time.time()
+            self._times[nick][func] = current_time
+            self._times[self.nick][func] = current_time
+            if not trigger.is_privmsg:
+                self._times[trigger.sender][func] = current_time
 
     def dispatch(self, pretrigger):
         args = pretrigger.args

sopel/loader.py

@@ -158,7 +158,9 @@ def clean_callable(func, config):
     func.unblockable = getattr(func, 'unblockable', False)
     func.priority = getattr(func, 'priority', 'medium')
     func.thread = getattr(func, 'thread', True)
-    func.rate = getattr(func, 'rate', 0)
+    func.user_rate = getattr(func, 'user_rate', 0)
+    func.channel_rate = getattr(func, 'channel_rate', 0)
+    func.global_rate = getattr(func, 'global_rate', 0)
 
     if not hasattr(func, 'event'):
         func.event = ['PRIVMSG']

sopel/module.py

@@ -248,8 +248,10 @@ def add_attribute(function):
     return add_attribute
 
 
-def rate(value):
-    """Decorate a function to limit how often a single user may trigger it.
+def rate(user_rate, channel_rate=0, global_rate=0):
+    """Decorate a function to limit how often it can be triggered. Without
+    arguments, the limit defaults to per-user. The global and channel parameters
+    are self-explanatory. A value of zero means no limit.
 
     If a function is given a rate of 20, a single user may only use that
     function once every 20 seconds. This limit applies to each user
@@ -260,7 +262,9 @@ def rate(value):
     threading.Timer() instead of sched, or rate limiting will not work properly.
     """
     def add_attribute(function):
-        function.rate = value
+        function.user_rate = user_rate
+        function.channel_rate = channel_rate
+        function.global_rate = global_rate
         return function
     return add_attribute