Fix Lack of Rate Limiting on login (API4:2019, Challenge 8)

api_views/users.py

@@ -1,4 +1,5 @@
 import re
+import time
 import jsonschema
 import jwt
 
@@ -16,6 +17,49 @@ def error_message_helper(msg):
         return '{ "status": "fail", "message": "' + msg + '"}'
 
 
+# --- Login rate limiting / lockout (Challenge 8: Lack of Resources & Rate Limiting) ---
+# In-memory tracker of recent failed login attempts, keyed by client IP + attempted
+# username so that repeated bad guesses against one account from one source get
+# throttled without needing an extra third-party dependency.
+_failed_login_attempts = {}
+LOGIN_MAX_ATTEMPTS = 5
+LOGIN_LOCKOUT_SECONDS = 60
+
+
+def _login_rate_limit_key(attempted_username):
+    client_ip = request.remote_addr or "unknown"
+    return "{}:{}".format(client_ip, attempted_username)
+
+
+def _login_lockout_remaining(rl_key):
+    """Returns seconds remaining in an active lockout for rl_key, or 0 if not locked out."""
+    entry = _failed_login_attempts.get(rl_key)
+    if not entry:
+        return 0
+    count, window_start = entry
+    elapsed = time.time() - window_start
+    if elapsed >= LOGIN_LOCKOUT_SECONDS:
+        # window has expired, forget the old attempts
+        _failed_login_attempts.pop(rl_key, None)
+        return 0
+    if count >= LOGIN_MAX_ATTEMPTS:
+        return max(1, int(LOGIN_LOCKOUT_SECONDS - elapsed))
+    return 0
+
+
+def _register_failed_login(rl_key):
+    now = time.time()
+    entry = _failed_login_attempts.get(rl_key)
+    if not entry or (now - entry[1]) >= LOGIN_LOCKOUT_SECONDS:
+        _failed_login_attempts[rl_key] = [1, now]
+    else:
+        entry[0] += 1
+
+
+def _clear_failed_logins(rl_key):
+    _failed_login_attempts.pop(rl_key, None)
+
+
 def get_all_users():
     return_value = jsonify({'users': User.get_all_users()})
     return return_value
@@ -85,19 +129,30 @@ def register_user():
 def login_user():
     request_data = request.get_json()
 
+    rl_key = _login_rate_limit_key((request_data or {}).get('username'))
+    lockout_remaining = _login_lockout_remaining(rl_key)
+    if lockout_remaining:
+        response = Response(
+            error_message_helper("Too many failed login attempts. Please try again later."),
+            429, mimetype="application/json")
+        response.headers['Retry-After'] = str(lockout_remaining)
+        return response
+
     try:
         # validate the data are in the correct form
         jsonschema.validate(request_data, login_user_schema)
         # fetching user data if the user exists
         user = User.query.filter_by(username=request_data.get('username')).first()
         if user and request_data.get('password') == user.password:
+            _clear_failed_logins(rl_key)
             auth_token = user.encode_auth_token(user.username)
             responseObject = {
                 'status': 'success',
                 'message': 'Successfully logged in.',
                 'auth_token': auth_token
             }
             return Response(json.dumps(responseObject), 200, mimetype="application/json")
+        _register_failed_login(rl_key)
         if vuln:  # Password Enumeration
             if user and request_data.get('password') != user.password:
                 return Response(error_message_helper("Password is not correct for the given username."), 200,