aspects/ratelimitaspects.py

   1 #!/usr/bin/python
   2 #-*- coding: utf-8 -*-
   3 #
   4 # S.Çağlar Onur <caglar@cs.princeton.edu>
   5
   6 from PLC.Config import Config
   7 from PLC.Faults import PLCPermissionDenied
   8
   9 from datetime import datetime, timedelta
  10
  11 from pyaspects.meta import MetaAspect
  12
  13 import memcache
  14
  15 class BaseRateLimit(object):
  16
  17     def __init__(self):
  18         self.config = Config("/etc/planetlab/plc_config")
  19
  20         # FIXME: change with Config values
  21         self.prefix = "ratelimit"
  22         self.minutes = 5 # The time period
  23         self.requests = 50 # Number of allowed requests in that time period
  24         self.expire_after = (self.minutes + 1) * 60
  25
  26         self.whitelist = []
  27
  28     def log(self, line):
  29         log = open("/var/log/plc_api_ratelimit.log", "a")
  30         date = datetime.now().strftime("%d/%m/%y %H:%M")
  31         log.write("%s - %s\n" % (date, line))
  32         log.flush()
  33
  34     def before(self, wobj, data, *args, **kwargs):
  35         # ratelimit_128.112.139.115_201011091532 = 1
  36         # ratelimit_128.112.139.115_201011091533 = 14
  37         # ratelimit_128.112.139.115_201011091534 = 11
  38         # Now, on every request we work out the keys for the past five minutes and use get_multi to retrieve them.
  39         # If the sum of those counters exceeds the maximum allowed for that time period, we block the request.
  40
  41         api_method_name = wobj.name
  42         api_method_source = wobj.source
  43
  44         try:
  45             api_method = args[0]["AuthMethod"]
  46         except:
  47             return
  48
  49         if api_method == "session":
  50             api_method_caller = args[0]["session"]
  51         elif api_method == "password" or api_method == "capability":
  52             api_method_caller = args[0]["Username"]
  53         elif api_method == "gpg":
  54             api_method_caller = args[0]["name"]
  55         elif api_method == "hmac" or api_method == "hmac_dummybox":
  56             api_method_caller = args[0]["node_id"]
  57         elif api_method == "anonymous":
  58             api_method_caller = "anonymous"
  59         else:
  60             api_method_caller = "unknown"
  61
  62         if api_method_source == None or api_method_source[0] == self.config.PLC_API_IP or api_method_source[0] in self.whitelist:
  63             return
  64
  65         if api_method_caller == None:
  66             self.log("%s called from %s with Username = None?" % (api_method_name, api_method_source[0]))
  67             return
  68
  69         mc = memcache.Client(["%s:11211" % self.config.PLC_API_HOST])
  70         now = datetime.now()
  71         current_key = "%s_%s_%s_%s" % (self.prefix, api_method_caller, api_method_source[0], now.strftime("%Y%m%d%H%M"))
  72
  73         keys_to_check = ["%s_%s_%s_%s" % (self.prefix, api_method_caller, api_method_source[0], (now - timedelta(minutes = minute)).strftime("%Y%m%d%H%M")) for minute in range(self.minutes + 1)]
  74
  75         try:
  76             value = mc.incr(current_key)
  77         except ValueError:
  78             value = None
  79
  80         if value == None:
  81             mc.set(current_key, 1, time=self.expire_after)
  82
  83         result = mc.get_multi(keys_to_check)
  84         total_requests = 0
  85         for i in result:
  86             total_requests += result[i]
  87
  88         if total_requests > self.requests:
  89             self.log("%s - %s" % (api_method_source[0], api_method_caller))
  90             raise PLCPermissionDenied, "Maximum allowed number of API calls exceeded"
  91
  92     def after(self, wobj, data, *args, **kwargs):
  93         return
  94
  95 class RateLimitAspect_class(BaseRateLimit):
  96     __metaclass__ = MetaAspect
  97     name = "ratelimitaspect_class"
  98
  99     def __init__(self):
 100         BaseRateLimit.__init__(self)
 101
 102     def before(self, wobj, data, *args, **kwargs):
 103         BaseRateLimit.before(self, wobj, data, *args, **kwargs)
 104
 105     def after(self, wobj, data, *args, **kwargs):
 106         BaseRateLimit.after(self, wobj, data, *args, **kwargs)
 107
 108 RateLimitAspect = RateLimitAspect_class