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