diff --git a/plugins/core_sieve.py b/plugins/core_sieve.py index 50a84c4..7a76515 100755 --- a/plugins/core_sieve.py +++ b/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', []) diff --git a/util/bucket.py b/util/bucket.py new file mode 100644 index 0000000..3a4a699 --- /dev/null +++ b/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) \ No newline at end of file