aspects/ratelimitaspects.py

   1 from PLC.Config import Config
   2 from PLC.Faults import *
   3
   4 from datetime import datetime, timedelta
   5 from pyaspects.meta import MetaAspect
   6 import memcache
   7
   8 class BaseRateLimit(object):
   9
  10     def __init__(self):
  11         self.config = Config("/etc/planetlab/plc_config")
  12
  13         self.prefix = "ratelimit"
  14         self.minutes = 3 # The time period
  15         self.requests = 10 # Number of allowed requests in that time period
  16         self.expire_after = (self.minutes + 1) * 60
  17
  18     def before(self, wobj, data, *args, **kwargs):
  19         # ratelimit_128.112.139.115_201011091532 = 1
  20         # ratelimit_128.112.139.115_201011091533 = 14
  21         # ratelimit_128.112.139.115_201011091534 = 11
  22         # Now, on every request we work out the keys for the past five minutes and use get_multi to retrieve them.
  23         # If the sum of those counters exceeds the maximum allowed for that time period, we block the request.
  24
  25         api_method_name = wobj.name
  26         api_method_source = wobj.source
  27
  28         if api_method_source == None or api_method_source[0] == self.config.PLC_API_HOST:
  29             return
  30
  31         mc = memcache.Client(['%s:11211'] % self.config.PLC_API_HOST)
  32         now = datetime.now()
  33         current_key = '%s_%s_%s' % (self.prefix, api_method_source[0], now.strftime('%Y%m%d%H%M'))
  34
  35         keys_to_check = ['%s_%s_%s' % (self.prefix, api_method_source[0], (now - timedelta(minutes = minute)).strftime('%Y%m%d%H%M')) for minute in range(self.minutes + 1)]
  36
  37         try:
  38             mc.incr(current_key)
  39         except ValueError:
  40             mc.set(current_key, 1, time=self.expire_after)
  41
  42         result = mc.get_multi(keys_to_check)
  43         total_requests = 0
  44         for i in result:
  45             total_requests += result[i]
  46
  47         if total_requests > self.requests:
  48             raise PLCPermissionDenied, "Maximum allowed number of API calls exceeded"
  49
  50     def after(self, wobj, data, *args, **kwargs):
  51         return
  52
  53 class RateLimitAspect_class(BaseRateLimit):
  54     __metaclass__ = MetaAspect
  55     name = "ratelimitaspect_class"
  56
  57     def __init__(self):
  58         BaseRateLimit.__init__(self)
  59
  60     def before(self, wobj, data, *args, **kwargs):
  61         BaseRateLimit.before(self, wobj, data, *args, **kwargs)
  62
  63     def after(self, wobj, data, *args, **kwargs):
  64         BaseRateLimit.after(self, wobj, data, *args, **kwargs)
  65
  66 RateLimitAspect = RateLimitAspect_class