Add rate limit and maximum amount of api keys (#1788)

Co-authored-by: Adrià Casajús <adria.casajus@proton.ch>
This commit is contained in:
Adrià Casajús 2023-06-29 17:21:00 +02:00 committed by GitHub
parent fc205157a8
commit d62bff8e46
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 54 additions and 1 deletions

View file

@ -534,3 +534,4 @@ SKIP_MX_LOOKUP_ON_CHECK = False
DISABLE_RATE_LIMIT = "DISABLE_RATE_LIMIT" in os.environ
SUBSCRIPTION_CHANGE_WEBHOOK = os.environ.get("SUBSCRIPTION_CHANGE_WEBHOOK", None)
MAX_API_KEYS = int(os.environ.get("MAX_API_KEYS", 30))

View file

@ -3,9 +3,11 @@ from flask_login import login_required, current_user
from flask_wtf import FlaskForm
from wtforms import StringField, validators
from app import config
from app.dashboard.base import dashboard_bp
from app.dashboard.views.enter_sudo import sudo_required
from app.db import Session
from app.extensions import limiter
from app.models import ApiKey
from app.utils import CSRFValidationForm
@ -14,9 +16,32 @@ class NewApiKeyForm(FlaskForm):
name = StringField("Name", validators=[validators.DataRequired()])
def clean_up_unused_or_old_api_keys(user_id: int):
total_keys = ApiKey.filter_by(user_id=user_id).count()
# Remove oldest unused
for api_key in (
ApiKey.filter_by(user_id=user_id, last_used=None)
.order_by(ApiKey.created_at.asc())
.all()
):
Session.delete(api_key)
total_keys -= 1
if total_keys <= config.MAX_API_KEYS:
return
# Clean up oldest used
for api_key in (
ApiKey.filter_by(user_id=user_id).order_by(ApiKey.last_used.asc()).all()
):
Session.delete(api_key)
total_keys -= 1
if total_keys <= config.MAX_API_KEYS:
return
@dashboard_bp.route("/api_key", methods=["GET", "POST"])
@login_required
@sudo_required
@limiter.limit("10/hour")
def api_key():
api_keys = (
ApiKey.filter(ApiKey.user_id == current_user.id)
@ -50,6 +75,7 @@ def api_key():
elif request.form.get("form-name") == "create":
if new_api_key_form.validate():
clean_up_unused_or_old_api_keys(current_user.id)
new_api_key = ApiKey.create(
name=new_api_key_form.name.data, user_id=current_user.id
)

View file

@ -1,10 +1,13 @@
from time import time
import arrow
from flask import url_for
from app import config
from app.dashboard.views.api_key import clean_up_unused_or_old_api_keys
from app.db import Session
from app.models import User, ApiKey
from tests.utils import login
from tests.utils import login, create_new_user
def test_api_key_page_requires_password(flask_client):
@ -87,3 +90,26 @@ def test_delete_all_api_keys(flask_client):
assert (
ApiKey.filter(ApiKey.user_id == user_2.id).count() == 1
) # assert that user 2 still has 1 API key
def test_cleanup_api_keys():
user = create_new_user()
ApiKey.create(
user_id=user.id, name="used", last_used=arrow.utcnow().shift(days=-3), times=1
)
ApiKey.create(
user_id=user.id, name="keep 1", last_used=arrow.utcnow().shift(days=-2), times=1
)
ApiKey.create(
user_id=user.id, name="keep 2", last_used=arrow.utcnow().shift(days=-1), times=1
)
ApiKey.create(user_id=user.id, name="not used", last_used=None, times=1)
Session.flush()
old_max_api_keys = config.MAX_API_KEYS
config.MAX_API_KEYS = 2
clean_up_unused_or_old_api_keys(user.id)
keys = ApiKey.filter_by(user_id=user.id).all()
assert len(keys) == 2
assert keys[0].name.find("keep") == 0
assert keys[1].name.find("keep") == 0
config.MAX_API_KEYS = old_max_api_keys