From 94395535892c1e939d5c61cc4dfac9d5928af058 Mon Sep 17 00:00:00 2001
From: Drew Short <warrick@sothr.com>
Date: Wed, 30 Jan 2019 23:11:17 -0600
Subject: [PATCH] Add a rate limiter * rate limit the unauthenticated endpoints
 to 2 requests a minute

---
 app.py           | 38 ++++++++++++++++++++++++++++----------
 requirements.txt |  1 +
 2 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/app.py b/app.py
index 5a65ef9..686bd25 100644
--- a/app.py
+++ b/app.py
@@ -6,16 +6,28 @@ from urllib.parse import urlparse, urljoin
 
 import flask
 from flask import Flask, redirect, render_template, request, g, flash, url_for
-from flask_login import LoginManager, login_required, login_user, logout_user, \
-    UserMixin
+from flask_limiter import Limiter
+from flask_limiter.util import get_remote_address
+from flask_login import (LoginManager,
+                         login_required,
+                         login_user,
+                         logout_user,
+                         UserMixin)
 from flask_wtf import CSRFProtect
 from werkzeug.contrib.fixers import ProxyFix
 
-from db import get_db, get_registration_codes, add_registration_code, \
-    expire_registration_code, delete_registration_code, get_registered_users, \
-    add_registered_user
-from forms import RegistrationForm, LoginForm, RegistrationCodeForm, \
-    ExpireRegistrationCodeForm
+from db import (
+    get_db,
+    get_registration_codes,
+    add_registration_code,
+    expire_registration_code,
+    delete_registration_code,
+    get_registered_users,
+    add_registered_user)
+from forms import (RegistrationForm,
+                   LoginForm,
+                   RegistrationCodeForm,
+                   ExpireRegistrationCodeForm)
 from register_new_matrix_user import register_new_user
 
 csrf = CSRFProtect()
@@ -82,6 +94,7 @@ def create_app():
 
 app = create_app()
 app.wsgi_app = ProxyFix(app.wsgi_app, num_proxies=1)
+limiter = Limiter(app, key_func=get_remote_address)
 
 log.info("Bound reverse proxy wsgi app")
 
@@ -111,7 +124,8 @@ class User(UserMixin):
 def is_safe_url(target):
     ref_url = urlparse(request.host_url)
     test_url = urlparse(urljoin(request.host_url, target))
-    return test_url.scheme in ('http', 'https') and ref_url.netloc == test_url.netloc
+    return test_url.scheme in (
+        'http', 'https') and ref_url.netloc == test_url.netloc
 
 
 def get_successful_registration_redirect():
@@ -137,15 +151,18 @@ def index():
 
 
 @app.route('/register', methods=('GET', 'POST'))
+@limiter.limit("2/minute")
 def registration():
     form = RegistrationForm()
     if form.validate_on_submit():
         if app.config.get("MATRIX_HOMESERVER") is None:
-            flash("Matrix Homeserver Currently Unavailable. Please Try Again Later!")
+            flash(
+                "Matrix Homeserver Currently Unavailable. Please Try Again Later!")
             return render_template('register.html', form=form)
         else:
             if app.config.get("MATRIX_SHARED_SECRET") is None:
-                flash("Registration Configuration Is Invalid. Contact Administrator!")
+                flash(
+                    "Registration Configuration Is Invalid. Contact Administrator!")
                 return render_template('register.html', form=form)
             else:
                 response = register_new_user(
@@ -216,6 +233,7 @@ def admin_expire_registration_code():
 
 
 @app.route('/admin/login', methods=('GET', 'POST'))
+@limiter.limit("2/minute")
 def admin_login():
     form = LoginForm()
     if form.validate_on_submit():
diff --git a/requirements.txt b/requirements.txt
index e397b2b..d5fbbfa 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
 flask==1.0.2
 flask-wtf==0.14
 flask-login==0.4.1
+flask-limiter==1.0.1
 requests==2.21.0
 safe
\ No newline at end of file