prototype

plugins/core_sieve.py

@@ -1,8 +1,10 @@
-import re
+from util import hook, bucket
 
-from util import hook
-from fnmatch import fnmatch
+TOKENS = 10
+RESTORE_RATE = 2
+MESSAGE_COST = 5
 
+buckets = {}
 
 @hook.sieve
 def sieve_suite(bot, input, func, kind, args):
@@ -12,6 +14,22 @@ def sieve_suite(bot, input, func, kind, args):
         return None
 
     if kind == "command":
+        uid = (input.nick, input.chan)
+
+        if not uid in buckets:
+            _bucket = bucket.TokenBucket(TOKENS, RESTORE_RATE)
+            _bucket.consume(MESSAGE_COST)
+            buckets[uid] = _bucket
+            return input
+
+        _bucket = buckets[uid]
+        if _bucket.consume(MESSAGE_COST):
+            return input
+        else:
+            print "pong!"
+            return None
+
+
         disabled_commands = conn.config.get('disabled_commands', [])
         if input.trigger in disabled_commands:
             return None
@@ -33,8 +51,6 @@ def sieve_suite(bot, input, func, kind, args):
         args["permissions"] = ["adminonly"]
 
     if args.get('permissions', False):
-
-
         mask = input.mask.lower()
 
         allowed_permissions = args.get('permissions', [])

util/bucket.py

@@ -0,0 +1,40 @@
+from time import time
+
+
+class TokenBucket(object):
+    """An implementation of the token bucket algorithm.
+    
+    >>> bucket = TokenBucket(80, 0.5)
+    >>> print bucket.consume(10)
+    True
+    >>> print bucket.consume(90)
+    False
+    """
+    def __init__(self, tokens, fill_rate):
+        """tokens is the total tokens in the bucket. fill_rate is the
+        rate in tokens/second that the bucket will be refilled."""
+        self.capacity = float(tokens)
+        self._tokens = float(tokens)
+        self.fill_rate = float(fill_rate)
+        self.timestamp = time()
+
+    def consume(self, tokens):
+        """Consume tokens from the bucket. Returns True if there were
+        sufficient tokens otherwise False."""
+        if tokens <= self.tokens:
+            self._tokens -= tokens
+        else:
+            return False
+        return True
+
+    def refill(self):
+        self._tokens = self.capacity
+
+    def get_tokens(self):
+        now = time()
+        if self._tokens < self.capacity:
+            delta = self.fill_rate * (now - self.timestamp)
+            self._tokens = min(self.capacity, self._tokens + delta)
+        self.timestamp = now
+        return self._tokens
+    tokens = property(get_tokens)