diff --git a/api_views/users.py b/api_views/users.py index 172540a..b50509e 100644 --- a/api_views/users.py +++ b/api_views/users.py @@ -1,4 +1,6 @@ import re +from collections import defaultdict, deque +import time import jsonschema import jwt @@ -9,6 +11,11 @@ from app import vuln +_LOGIN_LIMIT = 5 +_LOGIN_WINDOW = 60 +_login_attempts = defaultdict(deque) + + def error_message_helper(msg): if isinstance(msg, dict): return '{ "status": "fail", "message": "' + msg['error'] + '"}' @@ -83,6 +90,15 @@ def register_user(): def login_user(): + now = time.monotonic() + client = request.remote_addr or 'unknown' + attempts = _login_attempts[client] + while attempts and now - attempts[0] >= _LOGIN_WINDOW: + attempts.popleft() + if len(attempts) >= _LOGIN_LIMIT: + return Response(error_message_helper("Too many login attempts. Please try again later."), 429, + mimetype="application/json") + attempts.append(now) request_data = request.get_json() try: diff --git a/tests/test_rate_limit.py b/tests/test_rate_limit.py new file mode 100644 index 0000000..9cbed2b --- /dev/null +++ b/tests/test_rate_limit.py @@ -0,0 +1,9 @@ +from pathlib import Path + + +def test_login_has_ip_rate_limit(): + source = Path("api_views/users.py").read_text() + section = source[source.index("def login_user"):source.index("def token_validator")] + assert "request.remote_addr" in section + assert "_LOGIN_LIMIT" in section + assert "429" in section