59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
"""In-memory rate limiting for login attempts."""
|
|
|
|
import time
|
|
from collections import defaultdict
|
|
from fastapi import HTTPException
|
|
|
|
from app.config import settings
|
|
|
|
# In-memory store: IP address -> {"attempts": int, "lockout_until": float}
|
|
_rate_limits: dict[str, dict[str, float]] = defaultdict(
|
|
lambda: {"attempts": 0, "lockout_until": 0.0}
|
|
)
|
|
|
|
|
|
def check_rate_limit(ip: str) -> None:
|
|
"""
|
|
Check if the IP is currently locked out.
|
|
Raises a 429 HTTPException if locked out.
|
|
"""
|
|
now = time.time()
|
|
record = _rate_limits[ip]
|
|
|
|
if record["lockout_until"] > now:
|
|
remaining = int(record["lockout_until"] - now)
|
|
raise HTTPException(
|
|
status_code=429,
|
|
detail=f"Too many failed attempts. Try again in {remaining} seconds.",
|
|
headers={"Retry-After": str(remaining)},
|
|
)
|
|
|
|
# If lockout has expired, reset
|
|
if record["lockout_until"] > 0 and record["lockout_until"] <= now:
|
|
record["attempts"] = 0
|
|
record["lockout_until"] = 0.0
|
|
|
|
|
|
def record_failed_attempt(ip: str) -> None:
|
|
"""
|
|
Record a failed login attempt for the IP.
|
|
Applies lockout if MAX_ATTEMPTS is reached.
|
|
"""
|
|
now = time.time()
|
|
record = _rate_limits[ip]
|
|
record["attempts"] += 1
|
|
|
|
if record["attempts"] >= settings.LOGIN_MAX_ATTEMPTS:
|
|
record["lockout_until"] = now + settings.LOGIN_LOCKOUT_SECONDS
|
|
remaining = settings.LOGIN_LOCKOUT_SECONDS
|
|
raise HTTPException(
|
|
status_code=429,
|
|
detail=f"Too many failed attempts. Try again in {remaining} seconds.",
|
|
headers={"Retry-After": str(remaining)},
|
|
)
|
|
|
|
|
|
def clear_attempts(ip: str) -> None:
|
|
"""Clear all failed attempts for the IP (e.g., after successful login)."""
|
|
if ip in _rate_limits:
|
|
_rate_limits[ip] = {"attempts": 0, "lockout_until": 0.0}
|