4 # S.Çağlar Onur <caglar@cs.princeton.edu>
6 from PLC.Config import Config
7 from PLC.Faults import PLCPermissionDenied
9 from datetime import datetime, timedelta
11 from pyaspects.meta import MetaAspect
15 class BaseRateLimit(object):
18 self.config = Config("/etc/planetlab/plc_config")
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
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))
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.
41 api_method_name = wobj.name
42 api_method_source = wobj.source
45 api_method = args[0]["AuthMethod"]
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"
60 api_method_caller = "unknown"
62 if api_method_source == None or api_method_source[0] == self.config.PLC_API_IP or api_method_source[0] in self.whitelist:
65 if api_method_caller == None:
66 self.log("%s called from %s with Username = None?" % (api_method_name, api_method_source[0]))
69 mc = memcache.Client(["%s:11211" % self.config.PLC_API_HOST])
71 current_key = "%s_%s_%s_%s" % (self.prefix, api_method_caller, api_method_source[0], now.strftime("%Y%m%d%H%M"))
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)]
76 value = mc.incr(current_key)
81 mc.set(current_key, 1, time=self.expire_after)
83 result = mc.get_multi(keys_to_check)
86 total_requests += result[i]
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"
92 def after(self, wobj, data, *args, **kwargs):
95 class RateLimitAspect_class(BaseRateLimit):
96 __metaclass__ = MetaAspect
97 name = "ratelimitaspect_class"
100 BaseRateLimit.__init__(self)
102 def before(self, wobj, data, *args, **kwargs):
103 BaseRateLimit.before(self, wobj, data, *args, **kwargs)
105 def after(self, wobj, data, *args, **kwargs):
106 BaseRateLimit.after(self, wobj, data, *args, **kwargs)
108 RateLimitAspect = RateLimitAspect_class