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 PLC.Nodes import Node, Nodes
  10 from PLC.Persons import Person, Persons
  11 from PLC.Sessions import Session, Sessions
  12
  13 from datetime import datetime, timedelta
  14
  15 from pyaspects.meta import MetaAspect
  16
  17 import memcache
  18
  19 import os
  20 import sys
  21
  22 class BaseRateLimit(object):
  23
  24     def __init__(self):
  25         self.config = Config("/etc/planetlab/plc_config")
  26
  27         # FIXME: change with Config values
  28         self.prefix = "ratelimit"
  29         self.minutes = 5 # The time period
  30         self.requests = 50 # Number of allowed requests in that time period
  31         self.expire_after = (self.minutes + 1) * 60
  32
  33         self.whitelist = []
  34
  35     def log(self, line):
  36         log = open("/var/log/plc_api_ratelimit.log", "a")
  37         date = datetime.now().strftime("%d/%m/%y %H:%M")
  38         log.write("%s - %s\n" % (date, line))
  39         log.flush()
  40
  41     def mail(self, to):
  42         sendmail = os.popen("/usr/sbin/sendmail -N never -t -f%s" % self.config.PLC_MAIL_SUPPORT_ADDRESS, "w")
  43
  44         subject = "[PLCAPI] Maximum allowed number of API calls exceeded"
  45
  46         header = {'from': "%s Support <%s>" % (self.config.PLC_NAME, self.config.PLC_MAIL_SUPPORT_ADDRESS),
  47                'to': "%s, %s" % (to, self.config.PLC_MAIL_SUPPORT_ADDRESS),
  48                'version': sys.version.split(" ")[0],
  49                'subject': subject}
  50
  51         body = "Maximum allowed number of API calls exceeded for the user %s within the last %s minutes." % (to, self.minutes)
  52
  53         # Write headers
  54         sendmail.write(
  55 """
  56 Content-type: text/plain
  57 From: %(from)s
  58 Reply-To: %(from)s
  59 To: %(to)s
  60 X-Mailer: Python/%(version)s
  61 Subject: %(subject)s
  62
  63 """.lstrip() % header)
  64
  65         # Write body
  66         sendmail.write(body)
  67         # Done
  68         sendmail.close()
  69
  70     def before(self, wobj, data, *args, **kwargs):
  71         # ratelimit_128.112.139.115_201011091532 = 1
  72         # ratelimit_128.112.139.115_201011091533 = 14
  73         # ratelimit_128.112.139.115_201011091534 = 11
  74         # Now, on every request we work out the keys for the past five minutes and use get_multi to retrieve them.
  75         # If the sum of those counters exceeds the maximum allowed for that time period, we block the request.
  76
  77         api_method_name = wobj.name
  78         api_method_source = wobj.source
  79
  80         try:
  81             api_method = args[0]["AuthMethod"]
  82         except:
  83             return
  84
  85         # decode api_method_caller
  86         if api_method == "session":
  87             api_method_caller = Sessions(wobj.api, {'session_id': args[0]["session"]})
  88             if api_method_caller == []:
  89                 return
  90             elif api_method_caller[0]["person_id"] != None:
  91                 api_method_caller = Persons(wobj.api, api_method_caller[0]["person_id"])[0]["email"]
  92             elif api_method_caller[0]["node_id"] != None:
  93                 api_method_caller = Nodes(wobj.api, api_method_caller[0]["node_id"])[0]["hostname"]
  94             else:
  95                 api_method_caller = args[0]["session"]
  96         elif api_method == "password" or api_method == "capability":
  97             api_method_caller = args[0]["Username"]
  98         elif api_method == "gpg":
  99             api_method_caller = args[0]["name"]
 100         elif api_method == "hmac" or api_method == "hmac_dummybox":
 101             api_method_caller = args[0]["node_id"]
 102         elif api_method == "anonymous":
 103             api_method_caller = "anonymous"
 104         else:
 105             api_method_caller = "unknown"
 106
 107         # excludes
 108         if api_method_source == None or api_method_source[0] == self.config.PLC_API_IP or api_method_source[0] in self.whitelist:
 109             return
 110
 111         # sanity check
 112         if api_method_caller == None:
 113             self.log("%s called from %s with Username = None?" % (api_method_name, api_method_source[0]))
 114             return
 115
 116         # normalize unicode string otherwise memcache throws an exception
 117         api_method_caller = str(api_method_caller)
 118
 119         mc = memcache.Client(["%s:11211" % self.config.PLC_API_HOST])
 120         now = datetime.now()
 121
 122         current_key = "%s_%s_%s_%s" % (self.prefix, api_method_caller, api_method_source[0], now.strftime("%Y%m%d%H%M"))
 123         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)]
 124
 125         try:
 126             mc.incr(current_key)
 127         except ValueError:
 128             mc.set(current_key, 1, time=self.expire_after)
 129
 130         results = mc.get_multi(keys_to_check)
 131         total_requests = 0
 132         for i in results:
 133             total_requests += results[i]
 134
 135         if total_requests > self.requests:
 136             self.log("%s - %s" % (api_method_source[0], api_method_caller))
 137
 138             caller_key = "%s_%s" % (self.prefix, api_method_caller)
 139             if mc.get(caller_key) == None:
 140                 mc.set(caller_key, 1, time = self.expire_after)
 141                 if (api_method == "session" and api_method_caller.__contains__("@")) or (api_method == "password" or api_method == "capability"):
 142                     self.mail(api_method_caller)
 143
 144             raise PLCPermissionDenied, "Maximum allowed number of API calls exceeded"
 145
 146     def after(self, wobj, data, *args, **kwargs):
 147         return
 148
 149 class RateLimitAspect_class(BaseRateLimit):
 150     __metaclass__ = MetaAspect
 151     name = "ratelimitaspect_class"
 152
 153     def __init__(self):
 154         BaseRateLimit.__init__(self)
 155
 156     def before(self, wobj, data, *args, **kwargs):
 157         BaseRateLimit.before(self, wobj, data, *args, **kwargs)
 158
 159     def after(self, wobj, data, *args, **kwargs):
 160         BaseRateLimit.after(self, wobj, data, *args, **kwargs)
 161
 162 RateLimitAspect = RateLimitAspect_class