From 2a525d30691f94fa66dcfcb350bd9e979260de6b Mon Sep 17 00:00:00 2001
From: =?utf8?q?S=2E=C3=87a=C4=9Flar=20Onur?= <caglar@cs.princeton.edu>
Date: Tue, 16 Nov 2010 14:59:16 -0500
Subject: [PATCH] merge ratelimiting support from memcache branch

---
 PLCAPI.spec                 |  8 +++-
 aspects/__init__.py         |  7 +---
 aspects/ratelimitaspects.py | 78 +++++++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+), 6 deletions(-)
 create mode 100644 aspects/ratelimitaspects.py

diff --git a/PLCAPI.spec b/PLCAPI.spec
index 98236bb..2eb1116 100644
--- a/PLCAPI.spec
+++ b/PLCAPI.spec
@@ -47,7 +47,9 @@ Requires: python-twisted-words
 Requires: python-twisted-web
 # ldap
 Requires: python-ldap
-
+# for memcache
+Requires: python-memcached
+Requires: memcached
 ### avoid having yum complain about updates, as stuff is moving around
 # plc.d/api
 Conflicts: MyPLC <= 4.3
@@ -121,6 +123,10 @@ install -D -m 755 omf/omf_slicemgr.py $RPM_BUILD_ROOT/usr/bin/omf_slicemgr.py
 install -D -m 755 omf/reset_xmpp_pubsub_nodes.py $RPM_BUILD_ROOT/usr/bin/reset_xmpp_pubsub_nodes.py
 mkdir -p $RPM_BUILD_ROOT/var/log/omf
 
+# Install ratelimit log
+touch $RPM_BUILD_ROOT/var/log/plc_api_ratelimit.log
+chown apache:apache $RPM_BUILD_ROOT/var/log/plc_api_ratelimit.log
+
 %clean
 rm -rf $RPM_BUILD_ROOT
 
diff --git a/aspects/__init__.py b/aspects/__init__.py
index 1df5a9f..cdab6d6 100644
--- a/aspects/__init__.py
+++ b/aspects/__init__.py
@@ -1,16 +1,13 @@
-
-
 from pyaspects.weaver import weave_class_method
 
 from PLC.Method import Method
 from aspects.omfaspects import OMFAspect
-
-
+from aspects.ratelimitaspects import RateLimitAspect
 
 def apply_omf_aspect():
     # track all PLC methods to add OMF hooks
     weave_class_method(OMFAspect(), Method, "__call__")
-
+    weave_class_method(RateLimitAspect(), Method, "__call__")
 
 def apply_debugger_aspect():
     # just log all method calls w/ their parameters
diff --git a/aspects/ratelimitaspects.py b/aspects/ratelimitaspects.py
new file mode 100644
index 0000000..78548c2
--- /dev/null
+++ b/aspects/ratelimitaspects.py
@@ -0,0 +1,78 @@
+#!/usr/bin/python
+#-*- coding: utf-8 -*-
+#
+# S.Çağlar Onur <caglar@cs.princeton.edu>
+
+from PLC.Config import Config
+from PLC.Faults import PLCPermissionDenied
+
+from datetime import datetime, timedelta
+
+from pyaspects.meta import MetaAspect
+
+import memcache
+
+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
+
+    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
+
+        if api_method_source == None or api_method_source[0] == self.config.PLC_API_IP:
+            return
+
+        mc = memcache.Client(["%s:11211" % self.config.PLC_API_HOST])
+        now = datetime.now()
+        current_key = "%s_%s_%s" % (self.prefix, api_method_source[0], now.strftime("%Y%m%d%H%M"))
+
+        keys_to_check = ["%s_%s_%s" % (self.prefix, api_method_source[0], (now - timedelta(minutes = minute)).strftime("%Y%m%d%H%M")) for minute in range(self.minutes + 1)]
+
+        try:
+            mc.incr(current_key)
+        except ValueError:
+            mc.set(current_key, 1, time=self.expire_after)
+
+        result = mc.get_multi(keys_to_check)
+        total_requests = 0
+        for i in result:
+            total_requests += result[i]
+
+        if total_requests > self.requests:
+            log = open("/var/log/plc_api_ratelimit.log", "a")
+            date = datetime.now().strftime("%d/%m/%y %H:%M")
+            log.write("%s - %s\n" % (date, api_method_source[0]))
+            log.flush()
+            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
-- 
2.43.0