Limiter implementation

core/admin/mailu/internal/__init__.py

@@ -1,4 +1,4 @@
-from flask_limiter import RateLimitExceeded
+from mailu.limiter import RateLimitExceeded
 
 from mailu import utils
 from flask import current_app as app
@@ -20,13 +20,4 @@ def rate_limit_handler(e):
     return response
 
 
-@utils.limiter.request_filter
-def whitelist_webmail():
-    try:
-        return flask.request.headers["Client-Ip"] ==\
-            app.config["HOST_WEBMAIL"]
-    except:
-        return False
-
-
 from mailu.internal.views import *

core/admin/mailu/internal/views/auth.py

@@ -7,18 +7,21 @@ import flask_login
 import base64
 
 
+
 @internal.route("/auth/email")
-@utils.limiter.limit(
-    lambda: app.config["AUTH_RATELIMIT"],
-    lambda: flask.request.headers["Client-Ip"]
-)
 def nginx_authentication():
     """ Main authentication endpoint for Nginx email server
     """
+    utils.limiter.check(flask.request.headers["Client-Ip"])
     headers = nginx.handle_authentication(flask.request.headers)
     response = flask.Response()
     for key, value in headers.items():
         response.headers[key] = str(value)
+    if ("Auth-Status" in headers) and (headers["Auth-Status"]=="OK"):
+        utils.limiter.reset(flask.request.headers["Client-Ip"])
+    else:
+        utils.limiter.hit(flask.request.headers["Client-Ip"])
+
     return response
 
 

core/admin/mailu/limiter.py

@@ -0,0 +1,45 @@
+import limits
+import limits.storage
+import limits.strategies
+import ipaddress
+
+class RateLimitExceeded(Exception):
+    pass
+
+class Limiter:
+
+    def __init__(self):
+        self.storage = None
+        self.limiter = None
+        self.rate = None
+        self.subnet = None
+
+    def init_app(self, app):
+        self.storage = limits.storage.storage_from_string(app.config["RATELIMIT_STORAGE_URL"])
+        self.limiter = limits.strategies.MovingWindowRateLimiter(self.storage)
+        self.rate = limits.parse(app.config["AUTH_RATELIMIT"])
+        self.subnet = ipaddress.ip_network(app.config["SUBNET"])
+
+    def check(self,clientip):
+        # TODO: activate this code if we have limits at webmail level
+        #if ipaddress.ip_address(clientip) in self.subnet:
+        #    # no limits for internal requests (e.g. from webmail)
+        #    return
+        if not self.limiter.test(self.rate,"client-ip",clientip):
+            raise RateLimitExceeded()
+
+    def hit(self,clientip):
+        # TODO: activate this code if we have limits at webmail level
+        #if ipaddress.ip_address(clientip) in self.subnet:
+        #    # no limits for internal requests (e.g. from webmail)
+        #    return
+        if not self.limiter.hit(self.rate,"client-ip",clientip):
+            raise RateLimitExceeded()
+
+    def reset(self,clientip):
+        # TODO: activate this code if we have limits at webmail level
+        #if ipaddress.ip_address(clientip) in self.subnet:
+        #    # no limits for internal requests (e.g. from webmail)
+        #    return
+        # limit reset is not supported by the rate limit library
+        pass

core/admin/mailu/utils.py

@@ -1,11 +1,10 @@
-from mailu import models
+from mailu import models, limiter
 
 import flask
 import flask_login
 import flask_script
 import flask_migrate
 import flask_babel
-import flask_limiter
 
 from werkzeug.contrib import fixers
 
@@ -20,10 +19,8 @@ def handle_needs_login():
         flask.url_for('ui.login', next=flask.request.endpoint)
     )
 
-
-# Request rate limitation
-limiter = flask_limiter.Limiter(key_func=lambda: current_user.username)
-
+# Rate limiter
+limiter = limiter.Limiter()
 
 # Application translation
 babel = flask_babel.Babel()

core/admin/requirements.txt

@@ -7,7 +7,7 @@ Flask-migrate
 Flask-script
 Flask-wtf
 Flask-debugtoolbar
-Flask-limiter
+limits
 redis
 WTForms-Components
 socrate