From 889386b4a6e5d1eda52fec3dc1aa6ce043d3b59b Mon Sep 17 00:00:00 2001
From: Michael Wyraz <michael@wyraz.de>
Date: Fri, 6 Dec 2019 09:35:21 +0100
Subject: [PATCH] Limiter implementation

---
 core/admin/mailu/internal/__init__.py   | 11 +-----
 core/admin/mailu/internal/views/auth.py | 11 +++---
 core/admin/mailu/limiter.py             | 45 +++++++++++++++++++++++++
 core/admin/mailu/utils.py               |  9 ++---
 core/admin/requirements.txt             |  2 +-
 5 files changed, 57 insertions(+), 21 deletions(-)
 create mode 100644 core/admin/mailu/limiter.py

diff --git a/core/admin/mailu/internal/__init__.py b/core/admin/mailu/internal/__init__.py
index cf0ea3f7..95f2f782 100644
--- a/core/admin/mailu/internal/__init__.py
+++ b/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 *
diff --git a/core/admin/mailu/internal/views/auth.py b/core/admin/mailu/internal/views/auth.py
index 83a63953..af1c552f 100644
--- a/core/admin/mailu/internal/views/auth.py
+++ b/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
 
 
diff --git a/core/admin/mailu/limiter.py b/core/admin/mailu/limiter.py
new file mode 100644
index 00000000..4df78d27
--- /dev/null
+++ b/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
diff --git a/core/admin/mailu/utils.py b/core/admin/mailu/utils.py
index b11b1689..df23b8e7 100644
--- a/core/admin/mailu/utils.py
+++ b/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()
diff --git a/core/admin/requirements.txt b/core/admin/requirements.txt
index c68130db..bc6e772f 100644
--- a/core/admin/requirements.txt
+++ b/core/admin/requirements.txt
@@ -7,7 +7,7 @@ Flask-migrate
 Flask-script
 Flask-wtf
 Flask-debugtoolbar
-Flask-limiter
+limits
 redis
 WTForms-Components
 socrate