99 lines
3.2 KiB
Python
99 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from django.core.cache import cache
|
|
from django.http import HttpRequest
|
|
from django.utils import timezone
|
|
|
|
|
|
def _client_ip(request: HttpRequest) -> str:
|
|
forwarded_for = request.META.get("HTTP_X_FORWARDED_FOR", "").strip()
|
|
if forwarded_for:
|
|
return forwarded_for.split(",")[0].strip()
|
|
return request.META.get("REMOTE_ADDR", "unknown").strip() or "unknown"
|
|
|
|
|
|
def _identity_for_user(request: HttpRequest) -> str:
|
|
user = getattr(request, "user", None)
|
|
if user and getattr(user, "is_authenticated", False):
|
|
return f"user:{user.pk}"
|
|
username = (request.POST.get("username", "") if request.method == "POST" else "").strip().lower()
|
|
return f"anon:{username or _client_ip(request)}"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RateLimitState:
|
|
is_blocked: bool
|
|
remaining_seconds: int
|
|
attempts: int
|
|
|
|
|
|
def _attempts_key(namespace: str, identity: str) -> str:
|
|
return f"core-rate-limit:{namespace}:{identity}:attempts"
|
|
|
|
|
|
def _block_key(namespace: str, identity: str) -> str:
|
|
return f"core-rate-limit:{namespace}:{identity}:block"
|
|
|
|
|
|
def is_rate_limited(
|
|
request: HttpRequest,
|
|
*,
|
|
namespace: str,
|
|
max_attempts: int,
|
|
window_seconds: int,
|
|
block_seconds: int,
|
|
) -> RateLimitState:
|
|
now = timezone.now().timestamp()
|
|
identity = _identity_for_user(request)
|
|
block_until = cache.get(_block_key(namespace, identity))
|
|
if block_until and isinstance(block_until, (int, float)) and block_until > now:
|
|
return RateLimitState(True, max(0, int(block_until - now)), 0)
|
|
|
|
if cache.get(_attempts_key(namespace, identity), 0) >= max_attempts:
|
|
cache.set(
|
|
_block_key(namespace, identity),
|
|
now + max(block_seconds, 1),
|
|
timeout=block_seconds,
|
|
)
|
|
cache.delete(_attempts_key(namespace, identity))
|
|
return RateLimitState(True, max(block_seconds, 1), 0)
|
|
|
|
return RateLimitState(False, 0, int(cache.get(_attempts_key(namespace, identity), 0)))
|
|
|
|
|
|
def register_rate_limit_failure(
|
|
request: HttpRequest,
|
|
*,
|
|
namespace: str,
|
|
max_attempts: int,
|
|
window_seconds: int,
|
|
block_seconds: int,
|
|
) -> RateLimitState:
|
|
remaining = is_rate_limited(
|
|
request,
|
|
namespace=namespace,
|
|
max_attempts=max_attempts,
|
|
window_seconds=window_seconds,
|
|
block_seconds=block_seconds,
|
|
)
|
|
if remaining.is_blocked:
|
|
return remaining
|
|
|
|
identity = _identity_for_user(request)
|
|
attempts = int(cache.get(_attempts_key(namespace, identity), 0)) + 1
|
|
if attempts >= max_attempts:
|
|
now = timezone.now().timestamp()
|
|
cache.set(_block_key(namespace, identity), now + max(block_seconds, 1), timeout=block_seconds)
|
|
cache.delete(_attempts_key(namespace, identity))
|
|
return RateLimitState(True, max(block_seconds, 1), attempts)
|
|
|
|
cache.set(_attempts_key(namespace, identity), attempts, timeout=max(window_seconds, 1))
|
|
return RateLimitState(False, 0, attempts)
|
|
|
|
|
|
def clear_rate_limit(request: HttpRequest, *, namespace: str) -> None:
|
|
identity = _identity_for_user(request)
|
|
cache.delete(_attempts_key(namespace, identity))
|
|
cache.delete(_block_key(namespace, identity))
|