Merge pull request #209 from SibrenVasse/rate_limit

Implement rate limiting

app/api/views/auth.py

@@ -1,9 +1,8 @@
-import random
-
 import facebook
 import google.oauth2.credentials
 import googleapiclient.discovery
-from flask import jsonify, request
+import random
+from flask import jsonify, request, g
 from flask_cors import cross_origin
 from itsdangerous import Signer
 
@@ -17,13 +16,16 @@ from app.email_utils import (
     send_email,
     render,
 )
-from app.extensions import db
+from app.extensions import db, limiter
 from app.log import LOG
 from app.models import User, ApiKey, SocialAuth, AccountActivation
 
 
 @api_bp.route("/auth/login", methods=["POST"])
 @cross_origin()
+@limiter.limit(
+    "10/minute", deduct_when=lambda r: hasattr(g, "deduct_limit") and g.deduct_limit
+)
 def auth_login():
     """
     Authenticate user
@@ -52,6 +54,8 @@ def auth_login():
     user = User.filter_by(email=email).first()
 
     if not user or not user.check_password(password):
+        # Trigger rate limiter
+        g.deduct_limit = True
         return jsonify(error="Email or password incorrect"), 400
     elif not user.activated:
         return jsonify(error="Account not activated"), 400
@@ -113,6 +117,9 @@ def auth_register():
 
 @api_bp.route("/auth/activate", methods=["POST"])
 @cross_origin()
+@limiter.limit(
+    "10/minute", deduct_when=lambda r: hasattr(g, "deduct_limit") and g.deduct_limit
+)
 def auth_activate():
     """
     User enters the activation code to confirm their account.
@@ -136,16 +143,22 @@ def auth_activate():
 
     # do not use a different message to avoid exposing existing email
     if not user or user.activated:
+        # Trigger rate limiter
+        g.deduct_limit = True
         return jsonify(error="Wrong email or code"), 400
 
     account_activation = AccountActivation.get_by(user_id=user.id)
     if not account_activation:
+        # Trigger rate limiter
+        g.deduct_limit = True
         return jsonify(error="Wrong email or code"), 400
 
     if account_activation.code != code:
         # decrement nb tries
         account_activation.tries -= 1
         db.session.commit()
+        # Trigger rate limiter
+        g.deduct_limit = True
 
         if account_activation.tries == 0:
             AccountActivation.delete(account_activation.id)

app/auth/views/activate.py

@@ -1,14 +1,17 @@
-from flask import request, redirect, url_for, flash, render_template
+from flask import request, redirect, url_for, flash, render_template, g
 from flask_login import login_user, current_user
 
 from app import email_utils
 from app.auth.base import auth_bp
-from app.extensions import db
+from app.extensions import db, limiter
 from app.log import LOG
 from app.models import ActivationCode
 
 
 @auth_bp.route("/activate", methods=["GET", "POST"])
+@limiter.limit(
+    "10/minute", deduct_when=lambda r: hasattr(g, "deduct_limit") and g.deduct_limit
+)
 def activate():
     if current_user.is_authenticated:
         return (
@@ -21,6 +24,8 @@ def activate():
     activation_code: ActivationCode = ActivationCode.get_by(code=code)
 
     if not activation_code:
+        # Trigger rate limiter
+        g.deduct_limit = True
         return (
             render_template(
                 "auth/activate.html", error="Activation code cannot be found"

app/auth/views/fido.py

@@ -9,6 +9,7 @@ from flask import (
     flash,
     session,
     make_response,
+    g,
 )
 from flask_login import login_user
 from flask_wtf import FlaskForm
@@ -17,7 +18,7 @@ from wtforms import HiddenField, validators, BooleanField
 from app.auth.base import auth_bp
 from app.config import MFA_USER_ID
 from app.config import RP_ID, URL
-from app.extensions import db
+from app.extensions import db, limiter
 from app.log import LOG
 from app.models import User, Fido, MfaBrowser
 
@@ -30,6 +31,9 @@ class FidoTokenForm(FlaskForm):
 
 
 @auth_bp.route("/fido", methods=["GET", "POST"])
+@limiter.limit(
+    "10/minute", deduct_when=lambda r: hasattr(g, "deduct_limit") and g.deduct_limit
+)
 def fido():
     # passed from login page
     user_id = session.get(MFA_USER_ID)
@@ -57,6 +61,9 @@ def fido():
             flash(f"Welcome back {user.name}!", "success")
             # Redirect user to correct page
             return redirect(next_url or url_for("dashboard.index"))
+        else:
+            # Trigger rate limiter
+            g.deduct_limit = True
 
     # Handling POST requests
     if fido_token_form.validate_on_submit():
@@ -89,6 +96,8 @@ def fido():
         except Exception as e:
             LOG.error(f"An error occurred in WebAuthn verification process: {e}")
             flash("Key verification failed.", "warning")
+            # Trigger rate limiter
+            g.deduct_limit = True
             auto_activate = False
         else:
             user.fido_sign_count = new_sign_count

app/auth/views/forgot_password.py

@@ -1,9 +1,10 @@
-from flask import request, render_template, redirect, url_for, flash
+from flask import request, render_template, redirect, url_for, flash, g
 from flask_wtf import FlaskForm
 from wtforms import StringField, validators
 
 from app.auth.base import auth_bp
 from app.dashboard.views.setting import send_reset_password_email
+from app.extensions import limiter
 from app.models import User
 
 
@@ -12,6 +13,9 @@ class ForgotPasswordForm(FlaskForm):
 
 
 @auth_bp.route("/forgot_password", methods=["GET", "POST"])
+@limiter.limit(
+    "10/minute", deduct_when=lambda r: hasattr(g, "deduct_limit") and g.deduct_limit
+)
 def forgot_password():
     form = ForgotPasswordForm(request.form)
 
@@ -28,4 +32,7 @@ def forgot_password():
             send_reset_password_email(user)
             return redirect(url_for("auth.forgot_password"))
 
+        # Trigger rate limiter
+        g.deduct_limit = True
+
     return render_template("auth/forgot_password.html", form=form)

app/auth/views/login.py

@@ -1,10 +1,11 @@
-from flask import request, render_template, redirect, url_for, flash
+from flask import request, render_template, redirect, url_for, flash, g
 from flask_login import current_user
 from flask_wtf import FlaskForm
 from wtforms import StringField, validators
 
 from app.auth.base import auth_bp
 from app.auth.views.login_utils import after_login
+from app.extensions import limiter
 from app.log import LOG
 from app.models import User
 
@@ -15,6 +16,9 @@ class LoginForm(FlaskForm):
 
 
 @auth_bp.route("/login", methods=["GET", "POST"])
+@limiter.limit(
+    "10/minute", deduct_when=lambda r: hasattr(g, "deduct_limit") and g.deduct_limit
+)
 def login():
     if current_user.is_authenticated:
         LOG.d("user is already authenticated, redirect to dashboard")
@@ -27,9 +31,10 @@ def login():
     if form.validate_on_submit():
         user = User.filter_by(email=form.email.data.strip().lower()).first()
 
-        if not user:
+        if not user or not user.check_password(form.password.data):
-            flash("Email or password incorrect", "error")
+            # Trigger rate limiter
-        elif not user.check_password(form.password.data):
+            g.deduct_limit = True
+            form.password.data = None
             flash("Email or password incorrect", "error")
         elif not user.activated:
             show_resend_activation = True

app/auth/views/mfa.py

@@ -7,6 +7,7 @@ from flask import (
     session,
     make_response,
     request,
+    g,
 )
 from flask_login import login_user
 from flask_wtf import FlaskForm
@@ -14,7 +15,7 @@ from wtforms import BooleanField, StringField, validators
 
 from app.auth.base import auth_bp
 from app.config import MFA_USER_ID, URL
-from app.extensions import db
+from app.extensions import db, limiter
 from app.models import User, MfaBrowser
 
 
@@ -26,6 +27,9 @@ class OtpTokenForm(FlaskForm):
 
 
 @auth_bp.route("/mfa", methods=["GET", "POST"])
+@limiter.limit(
+    "10/minute", deduct_when=lambda r: hasattr(g, "deduct_limit") and g.deduct_limit
+)
 def mfa():
     # passed from login page
     user_id = session.get(MFA_USER_ID)
@@ -51,6 +55,9 @@ def mfa():
             flash(f"Welcome back {user.name}!", "success")
             # Redirect user to correct page
             return redirect(next_url or url_for("dashboard.index"))
+        else:
+            # Trigger rate limiter
+            g.deduct_limit = True
 
     if otp_token_form.validate_on_submit():
         totp = pyotp.TOTP(user.otp_secret)
@@ -84,6 +91,8 @@ def mfa():
 
         else:
             flash("Incorrect token", "warning")
+            # Trigger rate limiter
+            g.deduct_limit = True
             otp_token_form.token.data = None
 
     return render_template(

app/auth/views/recovery.py

@@ -1,13 +1,12 @@
 import arrow
-import pyotp
+from flask import request, render_template, redirect, url_for, flash, session, g
-from flask import request, render_template, redirect, url_for, flash, session
 from flask_login import login_user
 from flask_wtf import FlaskForm
 from wtforms import StringField, validators
 
 from app.auth.base import auth_bp
 from app.config import MFA_USER_ID
-from app.extensions import db
+from app.extensions import db, limiter
 from app.log import LOG
 from app.models import User, RecoveryCode
 
@@ -17,6 +16,9 @@ class RecoveryForm(FlaskForm):
 
 
 @auth_bp.route("/recovery", methods=["GET", "POST"])
+@limiter.limit(
+    "10/minute", deduct_when=lambda r: hasattr(g, "deduct_limit") and g.deduct_limit
+)
 def recovery_route():
     # passed from login page
     user_id = session.get(MFA_USER_ID)
@@ -41,6 +43,8 @@ def recovery_route():
 
         if recovery_code:
             if recovery_code.used:
+                # Trigger rate limiter
+                g.deduct_limit = True
                 flash("Code already used", "error")
             else:
                 del session[MFA_USER_ID]
@@ -60,6 +64,8 @@ def recovery_route():
                     LOG.debug("redirect user to dashboard")
                     return redirect(url_for("dashboard.index"))
         else:
+            # Trigger rate limiter
+            g.deduct_limit = True
             flash("Incorrect code", "error")
 
     return render_template("auth/recovery.html", recovery_form=recovery_form)

app/auth/views/reset_password.py

@@ -1,10 +1,10 @@
-from flask import request, flash, render_template, redirect, url_for
+from flask import request, flash, render_template, redirect, url_for, g
 from flask_login import login_user
 from flask_wtf import FlaskForm
 from wtforms import StringField, validators
 
 from app.auth.base import auth_bp
-from app.extensions import db
+from app.extensions import db, limiter
 from app.models import ResetPasswordCode
 
 
@@ -15,6 +15,9 @@ class ResetPasswordForm(FlaskForm):
 
 
 @auth_bp.route("/reset_password", methods=["GET", "POST"])
+@limiter.limit(
+    "10/minute", deduct_when=lambda r: hasattr(g, "deduct_limit") and g.deduct_limit
+)
 def reset_password():
     form = ResetPasswordForm(request.form)
 
@@ -25,6 +28,8 @@ def reset_password():
     )
 
     if not reset_password_code:
+        # Trigger rate limiter
+        g.deduct_limit = True
         error = (
             "The reset password link can be used only once. "
             "Please request a new link to reset password."

app/extensions.py

@@ -1,9 +1,22 @@
+from flask import request
+from flask_limiter import Limiter
+from flask_limiter.util import get_remote_address
 from flask_login import LoginManager
 from flask_migrate import Migrate
 from flask_sqlalchemy import SQLAlchemy
 
-
 db = SQLAlchemy()
 login_manager = LoginManager()
 login_manager.session_protection = "strong"
 migrate = Migrate(db=db)
+
+# Setup rate limit facility
+limiter = Limiter(key_func=get_remote_address)
+
+
+@limiter.request_filter
+def ip_whitelist():
+    # Uncomment line to test rate limit in dev environment
+    # return False
+    # No limit for local development
+    return request.remote_addr == "127.0.0.1"

requirements.in

@@ -40,3 +40,4 @@ google-auth-httplib2
 python-gnupg
 webauthn
 pyspf
+Flask-Limiter

requirements.txt

@@ -37,12 +37,13 @@ flask-admin==1.5.3        # via -r requirements.in
 flask-cors==3.0.8         # via -r requirements.in
 flask-debugtoolbar==0.10.1  # via -r requirements.in
 flask-httpauth==3.3.0     # via flask-profiler
+flask-limiter==1.3.1      # via -r requirements.in
 flask-login==0.4.1        # via -r requirements.in
 flask-migrate==2.5.2      # via -r requirements.in
 flask-profiler==1.8.1     # via -r requirements.in
 flask-sqlalchemy==2.4.0   # via -r requirements.in, flask-migrate
 flask-wtf==0.14.2         # via -r requirements.in
-flask==1.0.3              # via -r requirements.in, flask-admin, flask-cors, flask-debugtoolbar, flask-httpauth, flask-login, flask-migrate, flask-profiler, flask-sqlalchemy, flask-wtf
+flask==1.0.3              # via -r requirements.in, flask-admin, flask-cors, flask-debugtoolbar, flask-httpauth, flask-limiter, flask-login, flask-migrate, flask-profiler, flask-sqlalchemy, flask-wtf
 future==0.18.2            # via webauthn
 google-api-python-client==1.7.11  # via -r requirements.in
 google-auth-httplib2==0.0.3  # via -r requirements.in, google-api-python-client
@@ -59,6 +60,7 @@ jedi==0.13.3              # via ipython
 jinja2==2.10.1            # via flask, yacron
 jmespath==0.9.4           # via boto3, botocore
 jwcrypto==0.6.0           # via -r requirements.in
+limits==1.5.1             # via flask-limiter
 mako==1.0.12              # via alembic
 markupsafe==1.1.1         # via jinja2, mako
 more-itertools==7.0.0     # via pytest
@@ -98,7 +100,7 @@ ruamel.yaml==0.15.97      # via strictyaml
 s3transfer==0.2.1         # via boto3
 sentry-sdk==0.14.1        # via -r requirements.in
 simplejson==3.17.0        # via flask-profiler
-six==1.12.0               # via bcrypt, cryptography, flask-cors, google-api-python-client, google-auth, packaging, pip-tools, prompt-toolkit, pyopenssl, pytest, python-dateutil, sqlalchemy-utils, traitlets, webauthn
+six==1.12.0               # via bcrypt, cryptography, flask-cors, flask-limiter, google-api-python-client, google-auth, limits, packaging, pip-tools, prompt-toolkit, pyopenssl, pytest, python-dateutil, sqlalchemy-utils, traitlets, webauthn
 sqlalchemy-utils==0.36.1  # via -r requirements.in
 sqlalchemy==1.3.12        # via alembic, flask-sqlalchemy, sqlalchemy-utils
 strictyaml==1.0.2         # via yacron

server.py

@@ -1,9 +1,8 @@
-import os
-import ssl
-
 import arrow
 import flask_profiler
+import os
 import sentry_sdk
+import ssl
 from flask import Flask, redirect, url_for, render_template, request, jsonify, flash
 from flask_admin import Admin
 from flask_cors import cross_origin
@@ -11,6 +10,7 @@ from flask_login import current_user
 from sentry_sdk.integrations.aiohttp import AioHttpIntegration
 from sentry_sdk.integrations.flask import FlaskIntegration
 from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration
+from werkzeug.middleware.proxy_fix import ProxyFix
 
 from app import paddle_utils
 from app.admin_model import SLModelView, SLAdminIndexView
@@ -32,13 +32,12 @@ from app.config import (
 from app.dashboard.base import dashboard_bp
 from app.developer.base import developer_bp
 from app.discover.base import discover_bp
-from app.extensions import db, login_manager, migrate
+from app.extensions import db, login_manager, migrate, limiter
 from app.jose_utils import get_jwk_key
 from app.log import LOG
 from app.models import (
     Client,
     User,
-    Fido,
     ClientUser,
     Alias,
     RedirectUri,
@@ -50,8 +49,6 @@ from app.models import (
     Directory,
     Mailbox,
     DeletedAlias,
-    Contact,
-    EmailLog,
     Referral,
     AliasMailbox,
     Notification,
@@ -76,6 +73,10 @@ os.environ["OAUTHLIB_INSECURE_TRANSPORT"] = "1"
 
 def create_app() -> Flask:
     app = Flask(__name__)
+    # SimpleLogin is deployed behind NGINX
+    app.wsgi_app = ProxyFix(app.wsgi_app, num_proxies=1)
+    limiter.init_app(app)
+
     app.url_map.strict_slashes = False
 
     app.config["SQLALCHEMY_DATABASE_URI"] = DB_URI
@@ -369,6 +370,14 @@ def setup_error_page(app):
         else:
             return render_template("error/403.html"), 403
 
+    @app.errorhandler(429)
+    def forbidden(e):
+        LOG.error("Client hit rate limit on path %s", request.path)
+        if request.path.startswith("/api/"):
+            return jsonify(error="Rate limit exceeded"), 429
+        else:
+            return render_template("error/429.html"), 429
+
     @app.errorhandler(404)
     def page_not_found(e):
         if request.path.startswith("/api/"):

templates/error/429.html

@@ -0,0 +1,15 @@
+{% extends "error.html" %}
+
+{% block error_name %}
+  429
+{% endblock %}
+
+{% block error_description %}
+  Whoa, slow down there, pardner!
+{% endblock %}
+
+{% block suggestion %}
+  <a class="btn btn-primary" href="/">
+    <i class="fe fe-home mr-2"></i>Home Page
+  </a>
+{% endblock %}