73 lines
2.7 KiB
Python
73 lines
2.7 KiB
Python
"""Small in-process failed-attempt limiter for credential endpoints.
|
|
|
|
Fleet Ops runs as a single API process per deployment, so an in-memory sliding window
|
|
is sufficient to blunt online password guessing (and the scrypt CPU amplification that
|
|
comes with it) without adding Redis. Only *failed* attempts count, so legitimate users
|
|
and the automated test suite are never throttled.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from collections import deque
|
|
|
|
|
|
class FailedAttemptLimiter:
|
|
def __init__(self, *, max_failures: int, window_seconds: float) -> None:
|
|
self.max_failures = max_failures
|
|
self.window_seconds = window_seconds
|
|
self._failures: dict[str, deque[float]] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def _prune(self, key: str, now: float) -> deque[float]:
|
|
bucket = self._failures.setdefault(key, deque())
|
|
cutoff = now - self.window_seconds
|
|
while bucket and bucket[0] <= cutoff:
|
|
bucket.popleft()
|
|
if not bucket:
|
|
self._failures.pop(key, None)
|
|
return bucket
|
|
|
|
def retry_after_seconds(self, key: str) -> int:
|
|
"""Return >0 seconds to wait when the key is currently blocked, else 0."""
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
bucket = self._prune(key, now)
|
|
if len(bucket) < self.max_failures:
|
|
return 0
|
|
return max(1, int(bucket[0] + self.window_seconds - now + 0.999))
|
|
|
|
def record_failure(self, key: str) -> None:
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
self._prune(key, now)
|
|
self._failures.setdefault(key, deque()).append(now)
|
|
|
|
def reset(self, key: str) -> None:
|
|
with self._lock:
|
|
self._failures.pop(key, None)
|
|
|
|
|
|
class SlidingWindowLimiter:
|
|
"""Thread-safe request limiter where every accepted request consumes capacity."""
|
|
|
|
def __init__(self, *, max_requests: int, window_seconds: float) -> None:
|
|
self.max_requests = max_requests
|
|
self.window_seconds = window_seconds
|
|
self._requests: dict[str, deque[float]] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def consume(self, key: str) -> int:
|
|
"""Record an accepted request, or return the seconds until capacity is available."""
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
bucket = self._requests.setdefault(key, deque())
|
|
cutoff = now - self.window_seconds
|
|
while bucket and bucket[0] <= cutoff:
|
|
bucket.popleft()
|
|
if len(bucket) >= self.max_requests:
|
|
return max(1, int(bucket[0] + self.window_seconds - now + 0.999))
|
|
bucket.append(now)
|
|
return 0
|