+++ /dev/null
-#-*- coding: utf-8 -*-
-#
-# S.Çağlar Onur <caglar@cs.princeton.edu>
-
-from PLC.Config import Config
-from PLC.Faults import PLCPermissionDenied
-
-from PLC.Nodes import Node, Nodes
-from PLC.Persons import Person, Persons
-from PLC.Sessions import Session, Sessions
-
-from datetime import datetime, timedelta
-
-from pyaspects.meta import MetaAspect
-
-import memcache
-
-import os
-import sys
-import socket
-
-class BaseRateLimit(object):
-
- def __init__(self):
- self.config = Config("/etc/planetlab/plc_config")
-
- # FIXME: change with Config values
- self.prefix = "ratelimit"
- self.minutes = 5 # The time period
- self.requests = 50 # Number of allowed requests in that time period
- self.expire_after = (self.minutes + 1) * 60
-
- self.whitelist = []
-
- def log(self, line):
- log = open("/var/log/plc_api_ratelimit.log", "a")
- date = datetime.now().strftime("%d/%m/%y %H:%M")
- log.write("%s - %s\n" % (date, line))
- log.flush()
-
- def mail(self, to):
- sendmail = os.popen("/usr/sbin/sendmail -N never -t -f%s" % self.config.PLC_MAIL_SUPPORT_ADDRESS, "w")
-
- subject = "[PLCAPI] Maximum allowed number of API calls exceeded"
-
- header = {'from': "%s Support <%s>" % (self.config.PLC_NAME, self.config.PLC_MAIL_SUPPORT_ADDRESS),
- 'to': "%s, %s" % (to, self.config.PLC_MAIL_SUPPORT_ADDRESS),
- 'version': sys.version.split(" ")[0],
- 'subject': subject}
-
- body = "Maximum allowed number of API calls exceeded for the user %s within the last %s minutes." % (to, self.minutes)
-
- # Write headers
- sendmail.write(
-"""
-Content-type: text/plain
-From: %(from)s
-Reply-To: %(from)s
-To: %(to)s
-X-Mailer: Python/%(version)s
-Subject: %(subject)s
-
-""".lstrip() % header)
-
- # Write body
- sendmail.write(body)
- # Done
- sendmail.close()
-
- def before(self, wobj, data, *args, **kwargs):
- # ratelimit_128.112.139.115_201011091532 = 1
- # ratelimit_128.112.139.115_201011091533 = 14
- # ratelimit_128.112.139.115_201011091534 = 11
- # Now, on every request we work out the keys for the past five minutes and use get_multi to retrieve them.
- # If the sum of those counters exceeds the maximum allowed for that time period, we block the request.
-
- api_method_name = wobj.name
- api_method_source = wobj.source
-
- try:
- api_method = args[0]["AuthMethod"]
- except:
- return
-
- # decode api_method_caller
- if api_method == "session":
- api_method_caller = Sessions(wobj.api, {'session_id': args[0]["session"]})
- if api_method_caller == []:
- return
- elif api_method_caller[0]["person_id"] != None:
- api_method_caller = Persons(wobj.api, api_method_caller[0]["person_id"])[0]["email"]
- elif api_method_caller[0]["node_id"] != None:
- api_method_caller = Nodes(wobj.api, api_method_caller[0]["node_id"])[0]["hostname"]
- else:
- api_method_caller = args[0]["session"]
- elif api_method == "password" or api_method == "capability":
- api_method_caller = args[0]["Username"]
- elif api_method == "gpg":
- api_method_caller = args[0]["name"]
- elif api_method == "hmac" or api_method == "hmac_dummybox":
- api_method_caller = args[0]["node_id"]
- elif api_method == "anonymous":
- api_method_caller = "anonymous"
- else:
- api_method_caller = "unknown"
-
- # excludes
- if api_method_source == None or api_method_source[0] == socket.gethostbyname(self.config.PLC_API_HOST) or api_method_source[0] in self.whitelist:
- return
-
- # sanity check
- if api_method_caller == None:
- self.log("%s called from %s with Username = None?" % (api_method_name, api_method_source[0]))
- return
-
- # normalize unicode string otherwise memcache throws an exception
- api_method_caller = str(api_method_caller)
-
- mc = memcache.Client(["%s:11211" % self.config.PLC_API_HOST])
- now = datetime.now()
-
- current_key = "%s_%s_%s_%s" % (self.prefix, api_method_caller, api_method_source[0], now.strftime("%Y%m%d%H%M"))
- 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)]
-
- try:
- value = mc.incr(current_key)
- except ValueError:
- value = None
-
- if value == None:
- mc.set(current_key, 1, time=self.expire_after)
-
- results = mc.get_multi(keys_to_check)
- total_requests = 0
- for i in results:
- total_requests += results[i]
-
- if total_requests > self.requests:
- self.log("%s - %s" % (api_method_source[0], api_method_caller))
-
- caller_key = "%s_%s" % (self.prefix, api_method_caller)
- if mc.get(caller_key) == None:
- mc.set(caller_key, 1, time = self.expire_after)
- if (api_method == "session" and api_method_caller.__contains__("@")) or (api_method == "password" or api_method == "capability"):
- self.mail(api_method_caller)
-
- raise PLCPermissionDenied("Maximum allowed number of API calls exceeded")
-
- def after(self, wobj, data, *args, **kwargs):
- return
-
-class RateLimitAspect_class(BaseRateLimit, metaclass=MetaAspect):
- name = "ratelimitaspect_class"
-
- def __init__(self):
- BaseRateLimit.__init__(self)
-
- def before(self, wobj, data, *args, **kwargs):
- BaseRateLimit.before(self, wobj, data, *args, **kwargs)
-
- def after(self, wobj, data, *args, **kwargs):
- BaseRateLimit.after(self, wobj, data, *args, **kwargs)
-
-RateLimitAspect = RateLimitAspect_class