From 6598c4d259ae7bb9a6a0055ea229f999b282bd63 Mon Sep 17 00:00:00 2001
From: Ozzieisaacs <ozzie.fernandez.isaacs@googlemail.com>
Date: Sun, 4 Sep 2022 19:47:04 +0200
Subject: [PATCH] Add rate limit for opds

---
 cps/__init__.py | 1 +
 cps/admin.py    | 1 +
 cps/main.py     | 8 +++++++-
 cps/opds.py     | 8 +++++++-
 cps/web.py      | 3 ---
 5 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/cps/__init__.py b/cps/__init__.py
index 04df6443..1b91e317 100644
--- a/cps/__init__.py
+++ b/cps/__init__.py
@@ -158,6 +158,7 @@ def create_app():
         cache_buster.init_cache_busting(app)
     log.info('Starting Calibre Web...')
     limiter.init_app(app)
+    # limiter.limit("2/minute")(parent)
     Principal(app)
     lm.init_app(app)
     app.secret_key = os.getenv('SECRET_KEY', config_sql.get_flask_session_key(ub.session))
diff --git a/cps/admin.py b/cps/admin.py
index daf7a8e0..9707651b 100755
--- a/cps/admin.py
+++ b/cps/admin.py
@@ -119,6 +119,7 @@ def before_request():
                                  'admin.simulatedbchange',
                                  'admin.db_configuration',
                                  'web.login',
+                                 'web.login_post',
                                  'web.logout',
                                  'admin.load_dialogtexts',
                                  'admin.ajax_pathchooser'):
diff --git a/cps/main.py b/cps/main.py
index d3591c06..c3d8dd01 100644
--- a/cps/main.py
+++ b/cps/main.py
@@ -18,9 +18,14 @@
 
 import sys
 
-from . import create_app
+from . import create_app, limiter
 from .jinjia import jinjia
 from .remotelogin import remotelogin
+from flask import request
+
+
+def request_username():
+    return request.authorization.username
 
 def main():
     app = create_app()
@@ -56,6 +61,7 @@ def main():
     app.register_blueprint(tasks)
     app.register_blueprint(web)
     app.register_blueprint(opds)
+    limiter.limit("10/minute",key_func=request_username)(opds)
     app.register_blueprint(jinjia)
     app.register_blueprint(about)
     app.register_blueprint(shelf)
diff --git a/cps/opds.py b/cps/opds.py
index 60dbd551..bf3d77b8 100644
--- a/cps/opds.py
+++ b/cps/opds.py
@@ -31,12 +31,13 @@ from sqlalchemy.sql.expression import func, text, or_, and_, true
 from sqlalchemy.exc import InvalidRequestError, OperationalError
 from werkzeug.security import check_password_hash
 
-from . import constants, logger, config, db, calibre_db, ub, services, isoLanguages
+from . import constants, logger, config, db, calibre_db, ub, services, isoLanguages, limiter
 from .helper import get_download_link, get_book_cover
 from .pagination import Pagination
 from .web import render_read_books
 from .usermanagement import load_user_from_request
 from flask_babel import gettext as _
+from flask_limiter import RateLimitExceeded
 
 opds = Blueprint('opds', __name__)
 
@@ -479,6 +480,10 @@ def feed_search(term):
 
 
 def check_auth(username, password):
+    try:
+        limiter.check()
+    except RateLimitExceeded:
+        return False
     try:
         username = username.encode('windows-1252')
     except UnicodeEncodeError:
@@ -486,6 +491,7 @@ def check_auth(username, password):
     user = ub.session.query(ub.User).filter(func.lower(ub.User.name) ==
                                             username.decode('utf-8').lower()).first()
     if bool(user and check_password_hash(str(user.password), password)):
+        [limiter.limiter.storage.clear(k.key) for k in limiter.current_limits]
         return True
     else:
         ip_address = request.headers.get('X-Forwarded-For', request.remote_addr)
diff --git a/cps/web.py b/cps/web.py
index 6efd45f2..6a4b03fa 100755
--- a/cps/web.py
+++ b/cps/web.py
@@ -1277,9 +1277,6 @@ def handle_login_user(user, remember, message, category):
     [limiter.limiter.storage.clear(k.key) for k in limiter.current_limits]
     return redirect_back(url_for("web.index"))
 
-def error_logi():
-    flash(_(u"Wait one minute"), category="error")
-    return render_login()
 
 def render_login():
     next_url = request.args.get('next', default=url_for("web.index"), type=str)