aspects/ratelimitaspects.py

   1 #-*- coding: utf-8 -*-
   2 #
   3 # S.Çağlar Onur <caglar@cs.princeton.edu>
   4
   5 from PLC.Config import Config
   6 from PLC.Faults import PLCPermissionDenied
   7
   8 from PLC.Nodes import Node, Nodes
   9 from PLC.Persons import Person, Persons
  10 from PLC.Sessions import Session, Sessions
  11
  12 from datetime import datetime, timedelta
  13
  14 from pyaspects.meta import MetaAspect
  15
  16 import memcache
  17
  18 import os
  19 import sys
  20 import socket
  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] == socket.gethostbyname(self.config.PLC_API_HOST) 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             value = mc.incr(current_key)
 127         except ValueError:
 128             value = None
 129
 130         if value == None:
 131             mc.set(current_key, 1, time=self.expire_after)
 132
 133         results = mc.get_multi(keys_to_check)
 134         total_requests = 0
 135         for i in results:
 136             total_requests += results[i]
 137
 138         if total_requests > self.requests:
 139             self.log("%s - %s" % (api_method_source[0], api_method_caller))
 140
 141             caller_key = "%s_%s" % (self.prefix, api_method_caller)
 142             if mc.get(caller_key) == None:
 143                 mc.set(caller_key, 1, time = self.expire_after)
 144                 if (api_method == "session" and api_method_caller.__contains__("@")) or (api_method == "password" or api_method == "capability"):
 145                     self.mail(api_method_caller)
 146
 147             raise PLCPermissionDenied("Maximum allowed number of API calls exceeded")
 148
 149     def after(self, wobj, data, *args, **kwargs):
 150         return
 151
 152 class RateLimitAspect_class(BaseRateLimit, metaclass=MetaAspect):
 153     name = "ratelimitaspect_class"
 154
 155     def __init__(self):
 156         BaseRateLimit.__init__(self)
 157
 158     def before(self, wobj, data, *args, **kwargs):
 159         BaseRateLimit.before(self, wobj, data, *args, **kwargs)
 160
 161     def after(self, wobj, data, *args, **kwargs):
 162         BaseRateLimit.after(self, wobj, data, *args, **kwargs)
 163
 164 RateLimitAspect = RateLimitAspect_class