From b27686ccab62bfa2494cc6527c9f1504ce9bbdbc Mon Sep 17 00:00:00 2001
From: Nicolas F <ovdev@fratti.ch>
Date: Sun, 16 Jun 2019 23:08:23 +0200
Subject: [PATCH] Add login endpoint rate limiting

This doesn't discriminate between failed logins and successful
logins, but only counts POST requests. The limit is set to 6 per
hour.
---
 nyaa/__init__.py      | 5 ++++-
 nyaa/extensions.py    | 3 +++
 nyaa/views/account.py | 4 +++-
 requirements.txt      | 1 +
 4 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/nyaa/__init__.py b/nyaa/__init__.py
index f8830d1..ea06ef2 100644
--- a/nyaa/__init__.py
+++ b/nyaa/__init__.py
@@ -6,7 +6,7 @@ import flask
 from flask_assets import Bundle  # noqa F401
 
 from nyaa.api_handler import api_blueprint
-from nyaa.extensions import assets, cache, db, fix_paginate, toolbar
+from nyaa.extensions import assets, cache, db, fix_paginate, limiter, toolbar
 from nyaa.template_utils import bp as template_utils_bp
 from nyaa.template_utils import caching_url_for
 from nyaa.utils import random_string
@@ -128,4 +128,7 @@ def create_app(config):
     # Cache
     cache.init_app(app, config=app.config)
 
+    # Rate Limiting
+    limiter.init_app(app)
+
     return app
diff --git a/nyaa/extensions.py b/nyaa/extensions.py
index 6281007..1c6498f 100644
--- a/nyaa/extensions.py
+++ b/nyaa/extensions.py
@@ -5,12 +5,15 @@ from flask.config import Config
 from flask_assets import Environment
 from flask_caching import Cache
 from flask_debugtoolbar import DebugToolbarExtension
+from flask_limiter import Limiter
+from flask_limiter.util import get_remote_address
 from flask_sqlalchemy import BaseQuery, Pagination, SQLAlchemy
 
 assets = Environment()
 db = SQLAlchemy()
 toolbar = DebugToolbarExtension()
 cache = Cache()
+limiter = Limiter(key_func=get_remote_address)
 
 
 class LimitedPagination(Pagination):
diff --git a/nyaa/views/account.py b/nyaa/views/account.py
index b6a4a9e..5da4337 100644
--- a/nyaa/views/account.py
+++ b/nyaa/views/account.py
@@ -6,7 +6,7 @@ from ipaddress import ip_address
 import flask
 
 from nyaa import email, forms, models
-from nyaa.extensions import db
+from nyaa.extensions import db, limiter
 from nyaa.utils import sha1_hash
 from nyaa.views.users import get_activation_link, get_password_reset_link, get_serializer
 
@@ -15,6 +15,8 @@ bp = flask.Blueprint('account', __name__)
 
 
 @bp.route('/login', methods=['GET', 'POST'])
+@limiter.limit('6/hour', methods=['POST'],
+               error_message="You've tried logging in too many times, try again in an hour.")
 def login():
     if flask.g.user:
         return flask.redirect(redirect_url())
diff --git a/requirements.txt b/requirements.txt
index 214354c..e9b7a5e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -52,3 +52,4 @@ webassets==0.12.1
 Werkzeug==0.15.5
 WTForms==2.2.1
 Flask-Caching==1.7.2
+Flask-Limiter==1.0.1