diff --git a/app/config.py b/app/config.py index a33746b8..c0f280ca 100644 --- a/app/config.py +++ b/app/config.py @@ -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)) diff --git a/app/dashboard/views/api_key.py b/app/dashboard/views/api_key.py index 67dfc9ca..aa380755 100644 --- a/app/dashboard/views/api_key.py +++ b/app/dashboard/views/api_key.py @@ -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 ) diff --git a/tests/dashboard/test_api_keys.py b/tests/dashboard/test_api_keys.py index d2c36326..c00cda42 100644 --- a/tests/dashboard/test_api_keys.py +++ b/tests/dashboard/test_api_keys.py @@ -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