65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
"""Login rate limit and one-time upload tickets for the clips site."""
|
|
import secrets
|
|
import time
|
|
from collections import defaultdict, deque
|
|
from collections.abc import Callable
|
|
|
|
|
|
class LoginLimiter:
|
|
"""At most `limit` failed logins per key (an IP) within `window` seconds."""
|
|
|
|
def __init__(self, limit: int = 5, window: float = 600, now: Callable[[], float] = time.monotonic):
|
|
self.limit, self.window, self.now = limit, window, now
|
|
self.fails: dict[str, deque] = defaultdict(deque)
|
|
self.swept = now()
|
|
|
|
def _recent(self, key: str) -> deque:
|
|
q, cutoff = self.fails[key], self.now() - self.window
|
|
while q and q[0] <= cutoff:
|
|
q.popleft()
|
|
return q
|
|
|
|
def blocked(self, key: str) -> bool:
|
|
if key not in self.fails: # don't grow the table for every key that's merely checked
|
|
return False
|
|
return len(self._recent(key)) >= self.limit
|
|
|
|
def fail(self, key: str) -> None:
|
|
self._sweep()
|
|
self._recent(key).append(self.now())
|
|
|
|
def _sweep(self) -> None:
|
|
# forget keys whose strikes all expired, so guesses with made-up usernames can't grow the table forever
|
|
if self.now() - self.swept < self.window:
|
|
return
|
|
self.swept = self.now()
|
|
for key in [k for k in self.fails if not self._recent(k)]:
|
|
del self.fails[key]
|
|
|
|
def reset(self, key: str) -> None:
|
|
self.fails.pop(key, None)
|
|
|
|
|
|
class TicketStore:
|
|
"""One-time tickets that stand in for a Zipline token on /upload, valid `ttl` seconds."""
|
|
|
|
def __init__(self, ttl: float = 3600, now: Callable[[], float] = time.monotonic):
|
|
self.ttl, self.now = ttl, now
|
|
self.tickets: dict[str, tuple[str, float]] = {}
|
|
|
|
def _prune(self) -> None:
|
|
now = self.now()
|
|
for ticket in [t for t, (_, expires) in self.tickets.items() if expires <= now]:
|
|
del self.tickets[ticket]
|
|
|
|
def issue(self, token: str) -> str:
|
|
self._prune()
|
|
ticket = secrets.token_urlsafe(32)
|
|
self.tickets[ticket] = (token, self.now() + self.ttl)
|
|
return ticket
|
|
|
|
def redeem(self, ticket: str) -> str | None:
|
|
self._prune()
|
|
entry = self.tickets.pop(ticket, None)
|
|
return entry[0] if entry else None
|