93 lines
3.2 KiB
Python
93 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from collections import deque
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timedelta, timezone
|
|
from threading import Lock
|
|
|
|
from app.core.config import get_settings
|
|
from app.core.exceptions import TooManyRequestsError
|
|
|
|
settings = get_settings()
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _AttemptBucket:
|
|
timestamps: deque[datetime] = field(default_factory=deque)
|
|
last_seen: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
|
|
|
|
|
class LoginRateLimiter:
|
|
def __init__(self, attempts: int, window_seconds: int, max_buckets: int) -> None:
|
|
self.attempts = attempts
|
|
self.window = timedelta(seconds=window_seconds)
|
|
self.max_buckets = max_buckets
|
|
self._lock = Lock()
|
|
self._buckets: dict[str, _AttemptBucket] = {}
|
|
|
|
def build_key(self, username: str, client_host: str) -> str:
|
|
material = f"{username.lower()}:{client_host}".encode("utf-8")
|
|
return hashlib.sha256(material).hexdigest()
|
|
|
|
def is_limited(self, key: str) -> bool:
|
|
with self._lock:
|
|
now = datetime.now(timezone.utc)
|
|
self._cleanup(now=now)
|
|
bucket = self._buckets.get(key)
|
|
if bucket is None:
|
|
return False
|
|
self._touch(bucket, now=now)
|
|
self._prune(bucket, now=now)
|
|
return len(bucket.timestamps) >= self.attempts
|
|
|
|
def record_failure(self, key: str) -> None:
|
|
with self._lock:
|
|
now = datetime.now(timezone.utc)
|
|
self._cleanup(now=now)
|
|
bucket = self._buckets.get(key)
|
|
if bucket is None:
|
|
if len(self._buckets) >= self.max_buckets:
|
|
self._evict_oldest_bucket()
|
|
bucket = _AttemptBucket()
|
|
self._buckets[key] = bucket
|
|
self._touch(bucket, now=now)
|
|
self._prune(bucket, now=now)
|
|
bucket.timestamps.append(now)
|
|
|
|
def reset(self, key: str) -> None:
|
|
with self._lock:
|
|
self._buckets.pop(key, None)
|
|
|
|
def raise_limit_exceeded(self) -> None:
|
|
raise TooManyRequestsError("Too many login attempts. Please try again later.")
|
|
|
|
def _prune(self, bucket: _AttemptBucket, *, now: datetime) -> None:
|
|
cutoff = now - self.window
|
|
while bucket.timestamps and bucket.timestamps[0] < cutoff:
|
|
bucket.timestamps.popleft()
|
|
|
|
@staticmethod
|
|
def _touch(bucket: _AttemptBucket, *, now: datetime) -> None:
|
|
bucket.last_seen = now
|
|
|
|
def _cleanup(self, *, now: datetime) -> None:
|
|
stale_keys: list[str] = []
|
|
for key, bucket in self._buckets.items():
|
|
self._prune(bucket, now=now)
|
|
if not bucket.timestamps and bucket.last_seen < now - self.window:
|
|
stale_keys.append(key)
|
|
for key in stale_keys:
|
|
self._buckets.pop(key, None)
|
|
|
|
def _evict_oldest_bucket(self) -> None:
|
|
oldest_key = min(self._buckets, key=lambda bucket_key: self._buckets[bucket_key].last_seen)
|
|
self._buckets.pop(oldest_key, None)
|
|
|
|
|
|
login_rate_limiter = LoginRateLimiter(
|
|
attempts=settings.auth_rate_limit_attempts,
|
|
window_seconds=settings.auth_rate_limit_window_seconds,
|
|
max_buckets=settings.auth_rate_limit_max_buckets,
|
|
)
|