AutoML / gateway_auth.py
zukhriddinai's picture
Deploy public authenticated AutoML gateway
084b23f verified
Raw
History Blame Contribute Delete
9.77 kB
"""Public-safe account orchestration and bounded login throttling."""
from __future__ import annotations
import threading
import time
from dataclasses import dataclass
from typing import Any, Callable
import auth_storage
INVALID_LOGIN_MESSAGE = "Invalid username or password."
THROTTLED_LOGIN_MESSAGE = "Too many login attempts. Try again later."
THROTTLED_SIGNUP_MESSAGE = "Too many account requests. Try again later."
_SIGNUP_THROTTLE_USERNAME = "__account_signup__"
@dataclass(frozen=True)
class AuthResult:
ok: bool
message: str
session_token: str = ""
username: str = ""
account_status: str = ""
@dataclass
class _AttemptState:
failures: list[float]
locked_until: float = 0.0
updated_at: float = 0.0
class LoginThrottle:
"""Process-local brute-force limiter keyed by account and requester."""
def __init__(
self,
*,
max_failures: int = 5,
window_seconds: float = 300,
lock_seconds: float = 600,
max_entries: int = 10_000,
clock: Callable[[], float] = time.monotonic,
) -> None:
self.max_failures = max(1, int(max_failures))
self.window_seconds = max(1.0, float(window_seconds))
self.lock_seconds = max(1.0, float(lock_seconds))
self.max_entries = max(100, int(max_entries))
self._clock = clock
self._states: dict[tuple[str, str], _AttemptState] = {}
self._lock = threading.RLock()
@staticmethod
def _key(username: str, requester: str) -> tuple[str, str]:
return (
auth_storage.normalize_username(username).lower(),
str(requester or "unknown").strip() or "unknown",
)
def _prune_state(self, state: _AttemptState, now: float) -> None:
cutoff = now - self.window_seconds
state.failures[:] = [stamp for stamp in state.failures if stamp >= cutoff]
if state.locked_until and now >= state.locked_until:
state.locked_until = 0.0
state.failures.clear()
state.updated_at = now
def _bound_entries(self) -> None:
excess = len(self._states) - self.max_entries
if excess <= 0:
return
oldest = sorted(self._states.items(), key=lambda item: item[1].updated_at)
for key, _state in oldest[:excess]:
self._states.pop(key, None)
def allowed(self, username: str, requester: str) -> bool:
now = self._clock()
key = self._key(username, requester)
with self._lock:
state = self._states.get(key)
if state is None:
return True
self._prune_state(state, now)
if not state.failures and not state.locked_until:
self._states.pop(key, None)
return True
return state.locked_until <= now
def record_failure(self, username: str, requester: str) -> None:
now = self._clock()
key = self._key(username, requester)
with self._lock:
state = self._states.setdefault(key, _AttemptState([]))
self._prune_state(state, now)
state.failures.append(now)
if len(state.failures) >= self.max_failures:
state.locked_until = now + self.lock_seconds
state.updated_at = now
self._bound_entries()
def record_success(self, username: str, requester: str) -> None:
with self._lock:
self._states.pop(self._key(username, requester), None)
def failure_count(self, username: str, requester: str) -> int:
now = self._clock()
with self._lock:
state = self._states.get(self._key(username, requester))
if state is None:
return 0
self._prune_state(state, now)
return len(state.failures)
DEFAULT_LOGIN_THROTTLE = LoginThrottle()
DEFAULT_SIGNUP_THROTTLE = LoginThrottle(
max_failures=5,
window_seconds=3600,
lock_seconds=3600,
)
def requester_key(request: Any | None) -> str:
"""Use the server-observed peer address, never a spoofable forwarded header."""
client = getattr(request, "client", None)
host = str(getattr(client, "host", "") or "").strip()
return host or "unknown"
def login(
username: str,
password: str,
requester: str,
*,
throttle: LoginThrottle = DEFAULT_LOGIN_THROTTLE,
) -> AuthResult:
normalized = auth_storage.normalize_username(username)
if not throttle.allowed(normalized, requester):
return AuthResult(False, THROTTLED_LOGIN_MESSAGE)
status = auth_storage.account_authentication_status(normalized, password)
if status == "invalid":
throttle.record_failure(normalized, requester)
return AuthResult(False, INVALID_LOGIN_MESSAGE)
throttle.record_success(normalized, requester)
if status == auth_storage.ACCOUNT_STATUS_PENDING:
return AuthResult(
False,
"Administrators have not approved this account yet.",
username=normalized,
account_status=status,
)
if status == auth_storage.ACCOUNT_STATUS_DENIED:
return AuthResult(
False,
"Your request for account creation was denied by administrators. "
"Choose whether to re-send or delete the request below.",
username=normalized,
account_status=status,
)
return AuthResult(
True,
f"Signed in as {normalized}.",
session_token=auth_storage.issue_session_token(normalized),
username=normalized,
account_status=auth_storage.ACCOUNT_STATUS_APPROVED,
)
def request_signup(
username: str,
password: str,
confirm_password: str,
requester: str,
*,
throttle: LoginThrottle = DEFAULT_SIGNUP_THROTTLE,
) -> AuthResult:
normalized = auth_storage.normalize_username(username)
if not auth_storage.signups_allowed():
return AuthResult(False, "Account creation is disabled for this deployment.")
if not throttle.allowed(_SIGNUP_THROTTLE_USERNAME, requester):
return AuthResult(False, THROTTLED_SIGNUP_MESSAGE)
if password != confirm_password:
return AuthResult(False, "Passwords do not match.")
try:
user = auth_storage.request_account_creation(normalized, password)
except ValueError as exc:
throttle.record_failure(_SIGNUP_THROTTLE_USERNAME, requester)
return AuthResult(False, f"Could not submit account request: {exc}")
except Exception:
throttle.record_failure(_SIGNUP_THROTTLE_USERNAME, requester)
return AuthResult(
False,
"Could not submit the account request. Try again later.",
)
throttle.record_failure(_SIGNUP_THROTTLE_USERNAME, requester)
return AuthResult(
False,
f"Request for account creation for {user['username']} has been sent to administrators.",
username=str(user["username"]),
account_status=auth_storage.ACCOUNT_STATUS_PENDING,
)
def handle_denied_request(
username: str,
password: str,
choice: str,
) -> AuthResult:
normalized = auth_storage.normalize_username(username)
try:
if choice == "Re-send request":
auth_storage.resubmit_denied_account(normalized, password)
return AuthResult(
False,
"Your account creation request has been re-sent to administrators.",
username=normalized,
account_status=auth_storage.ACCOUNT_STATUS_PENDING,
)
if choice == "Do not re-send; delete my account request":
auth_storage.delete_denied_account(normalized, password)
return AuthResult(
False,
"Your denied account request and all related database information "
"have been deleted.",
)
return AuthResult(
False,
"Choose whether to re-send or delete the account request.",
username=normalized,
account_status=auth_storage.ACCOUNT_STATUS_DENIED,
)
except (PermissionError, ValueError):
return AuthResult(
False,
"Could not process the denied account request. Re-enter the password "
"used to sign up and try again.",
username=normalized,
account_status=auth_storage.ACCOUNT_STATUS_DENIED,
)
def change_password(
username: str,
old_password: str,
new_password: str,
confirm_password: str,
) -> AuthResult:
normalized = auth_storage.normalize_username(username)
if not normalized or not old_password or not new_password or not confirm_password:
return AuthResult(False, "Enter your username, current password, and the new password twice.")
if new_password != confirm_password:
return AuthResult(False, "New passwords do not match.")
if old_password == new_password:
return AuthResult(False, "Choose a new password that differs from your current password.")
try:
auth_storage.update_password(normalized, old_password, new_password)
except PermissionError:
return AuthResult(
False,
"Could not update password: Invalid username or current password.",
)
except ValueError as exc:
return AuthResult(False, f"Could not update password: {exc}")
return AuthResult(
True,
"Password updated. Sign in again with your new password.",
username=normalized,
)
__all__ = [
"AuthResult",
"LoginThrottle",
"DEFAULT_LOGIN_THROTTLE",
"DEFAULT_SIGNUP_THROTTLE",
"requester_key",
"login",
"request_signup",
"handle_denied_request",
"change_password",
]