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         # FIXME: change with Config values
  14         self.prefix = "ratelimit"
  15         self.minutes = 5 # The time period
  16         self.requests = 50 # Number of allowed requests in that time period
  17         self.expire_after = (self.minutes + 1) * 60
  18
  19     def before(self, wobj, data, *args, **kwargs):
  20         # ratelimit_128.112.139.115_201011091532 = 1
  21         # ratelimit_128.112.139.115_201011091533 = 14
  22         # ratelimit_128.112.139.115_201011091534 = 11
  23         # Now, on every request we work out the keys for the past five minutes and use get_multi to retrieve them.
  24         # If the sum of those counters exceeds the maximum allowed for that time period, we block the request.
  25
  26         api_method_name = wobj.name
  27         api_method_source = wobj.source
  28
  29         if api_method_source == None or api_method_source[0] == self.config.PLC_API_IP:
  30             return
  31
  32         mc = memcache.Client(["%s:11211" % self.config.PLC_API_HOST])
  33         now = datetime.now()
  34         current_key = "%s_%s_%s" % (self.prefix, api_method_source[0], now.strftime("%Y%m%d%H%M"))
  35
  36         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)]
  37
  38         try:
  39             mc.incr(current_key)
  40         except ValueError:
  41             mc.set(current_key, 1, time=self.expire_after)
  42
  43         result = mc.get_multi(keys_to_check)
  44         total_requests = 0
  45         for i in result:
  46             total_requests += result[i]
  47
  48         if total_requests > self.requests:
  49             log = open("/var/log/api_ratelimit.log", "a")
  50             date = datetime.now().strftime("%d/%m/%y %H:%M")
  51             log.write("%s - %s\n" % (date, api_method_source[0]))
  52             log.flush()
  53             raise PLCPermissionDenied, "Maximum allowed number of API calls exceeded"
  54
  55     def after(self, wobj, data, *args, **kwargs):
  56         return
  57
  58 class RateLimitAspect_class(BaseRateLimit):
  59     __metaclass__ = MetaAspect
  60     name = "ratelimitaspect_class"
  61
  62     def __init__(self):
  63         BaseRateLimit.__init__(self)
  64
  65     def before(self, wobj, data, *args, **kwargs):
  66         BaseRateLimit.before(self, wobj, data, *args, **kwargs)
  67
  68     def after(self, wobj, data, *args, **kwargs):
  69         BaseRateLimit.after(self, wobj, data, *args, **kwargs)
  70
  71 RateLimitAspect = RateLimitAspect_class