Add rate limiting to VAmPI login

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:

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