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