112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
|
import logging
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
from enum import StrEnum
|
||
|
|
from typing import NamedTuple
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
AUTH_FAIL_LIMIT = 10
|
||
|
|
AUTH_LOCKOUT_SECONDS = 300
|
||
|
|
AUTH_RATE_SCOPE_BUCKETS = 10_000
|
||
|
|
|
||
|
|
|
||
|
|
class AuthRateScope(StrEnum):
|
||
|
|
TOKEN = "token"
|
||
|
|
PASSWORD = "password"
|
||
|
|
DEVICE_TOKEN = "device_token"
|
||
|
|
BOOTSTRAP = "bootstrap"
|
||
|
|
|
||
|
|
|
||
|
|
class _ScopeState(NamedTuple):
|
||
|
|
failures: int
|
||
|
|
locked_until: float
|
||
|
|
|
||
|
|
|
||
|
|
class AuthRateLimiter:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
fail_limit: int = AUTH_FAIL_LIMIT,
|
||
|
|
lockout_seconds: float = AUTH_LOCKOUT_SECONDS,
|
||
|
|
max_buckets: int = AUTH_RATE_SCOPE_BUCKETS,
|
||
|
|
):
|
||
|
|
self._fail_limit = fail_limit
|
||
|
|
self._lockout_seconds = lockout_seconds
|
||
|
|
self._max_buckets = max_buckets
|
||
|
|
self._buckets: dict[tuple[str, AuthRateScope], _ScopeState] = {}
|
||
|
|
self._lock = threading.Lock()
|
||
|
|
|
||
|
|
def _cleanup_expired(self) -> None:
|
||
|
|
if len(self._buckets) <= self._max_buckets:
|
||
|
|
return
|
||
|
|
now = time.monotonic()
|
||
|
|
expired = [
|
||
|
|
key
|
||
|
|
for key, state in self._buckets.items()
|
||
|
|
if now >= state.locked_until and state.failures < self._fail_limit
|
||
|
|
]
|
||
|
|
for key in expired:
|
||
|
|
del self._buckets[key]
|
||
|
|
|
||
|
|
def record_failure(self, identifier: str, scope: AuthRateScope) -> bool:
|
||
|
|
with self._lock:
|
||
|
|
self._cleanup_expired()
|
||
|
|
key = (identifier, scope)
|
||
|
|
now = time.monotonic()
|
||
|
|
state = self._buckets.get(key)
|
||
|
|
|
||
|
|
if state is not None and now < state.locked_until:
|
||
|
|
return False
|
||
|
|
|
||
|
|
new_failures = (state.failures + 1) if state else 1
|
||
|
|
locked_until = now + self._lockout_seconds if new_failures >= self._fail_limit else 0.0
|
||
|
|
self._buckets[key] = _ScopeState(new_failures, locked_until)
|
||
|
|
|
||
|
|
if new_failures >= self._fail_limit:
|
||
|
|
logger.warning(
|
||
|
|
"Auth rate limit locked: identifier=%s scope=%s",
|
||
|
|
identifier,
|
||
|
|
scope.value,
|
||
|
|
)
|
||
|
|
|
||
|
|
if len(self._buckets) > self._max_buckets * 2:
|
||
|
|
self._force_evict(now)
|
||
|
|
|
||
|
|
return True
|
||
|
|
|
||
|
|
def reset(self, identifier: str, scope: AuthRateScope) -> None:
|
||
|
|
key = (identifier, scope)
|
||
|
|
with self._lock:
|
||
|
|
self._buckets.pop(key, None)
|
||
|
|
|
||
|
|
def is_locked(self, identifier: str, scope: AuthRateScope) -> bool:
|
||
|
|
key = (identifier, scope)
|
||
|
|
with self._lock:
|
||
|
|
state = self._buckets.get(key)
|
||
|
|
if state is None:
|
||
|
|
return False
|
||
|
|
return time.monotonic() < state.locked_until
|
||
|
|
|
||
|
|
def remaining_attempts(self, identifier: str, scope: AuthRateScope) -> int:
|
||
|
|
key = (identifier, scope)
|
||
|
|
with self._lock:
|
||
|
|
state = self._buckets.get(key)
|
||
|
|
if state is None:
|
||
|
|
return self._fail_limit
|
||
|
|
now = time.monotonic()
|
||
|
|
if state.locked_until > 0.0 and now >= state.locked_until:
|
||
|
|
return self._fail_limit
|
||
|
|
return max(0, self._fail_limit - state.failures)
|
||
|
|
|
||
|
|
def _force_evict(self, now: float) -> None:
|
||
|
|
by_expiry = sorted(
|
||
|
|
self._buckets.items(),
|
||
|
|
key=lambda item: item[1].locked_until,
|
||
|
|
)
|
||
|
|
to_remove = len(self._buckets) - self._max_buckets
|
||
|
|
for i in range(to_remove):
|
||
|
|
del self._buckets[by_expiry[i][0]]
|
||
|
|
|
||
|
|
|
||
|
|
auth_rate_limiter = AuthRateLimiter()
|