webmail-client/backend/app/identity/service.py

"""Registration, authentication, and login rate limiting (SPEC §6.1)."""

from __future__ import annotations

import uuid
from collections.abc import Awaitable, Callable
from dataclasses import dataclass

from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import get_settings
from app.identity.handles import InvalidHandle, full_handle, normalize_localpart
from app.identity.passwords import hash_password, verify_password
from app.models import AppSettings, User
from app.redis_client import get_redis

# SPEC §6.1: 10 attempts / 15 min per handle AND per IP.
_RL_MAX = 10
_RL_WINDOW = 15 * 60


class RegistrationError(Exception):
    """Registration failed for a user-actionable reason."""

    def __init__(self, code: str, message: str) -> None:
        super().__init__(message)
        self.code = code
        self.message = message


@dataclass(frozen=True)
class RegistrationStatus:
    """What the signup page may offer right now."""

    open: bool
    # True when the instance has no users: the next signup creates the administrator.
    bootstrap: bool


@dataclass(frozen=True)
class RateLimited(Exception):
    """Raised when the login rate limit is exceeded."""

    retry_after: int


class TotpRequired(Exception):
    """Password verified but the account has 2FA enabled and no code was given."""


class TotpInvalid(Exception):
    """Password verified but the supplied TOTP code was wrong."""


class AccountBlocked(Exception):
    """Password verified but an administrator has blocked the account."""


def resolve_handle(handle_input: str) -> str:
    """Normalize a login identifier to the full ``localpart@APP_DOMAIN`` form."""
    text = handle_input.strip().lower()
    if "@" in text:
        return text
    return f"{normalize_localpart(text)}@{get_settings().app_domain}"


async def app_settings(db: AsyncSession, *, for_update: bool = False) -> AppSettings:
    """Load the singleton settings row, creating it if an old database lacks one.

    ``for_update`` takes a row lock, which is what serialises concurrent signups.
    A row conjured here defaults to *closed*: the safe answer when we cannot tell
    what the operator intended.
    """
    stmt = select(AppSettings).where(AppSettings.id == 1)
    row = await db.scalar(stmt.with_for_update() if for_update else stmt)
    if row is None:
        await db.execute(
            pg_insert(AppSettings).values(id=1, registration_open=False).on_conflict_do_nothing()
        )
        row = await db.scalar(stmt.with_for_update() if for_update else stmt)
        assert row is not None  # just inserted in this transaction
    return row


async def _has_users(db: AsyncSession) -> bool:
    return await db.scalar(select(User.id).limit(1)) is not None


async def registration_status(db: AsyncSession) -> RegistrationStatus:
    """Whether signup is currently accepted, and whether it would bootstrap an admin."""
    bootstrap = not await _has_users(db)
    settings = await app_settings(db)
    return RegistrationStatus(open=bootstrap or settings.registration_open, bootstrap=bootstrap)


async def register_user(
    db: AsyncSession, *, localpart: str, password: str, recovery_email: str | None
) -> User:
    """Create a new user. Raises :class:`RegistrationError` on validation/uniqueness.

    The first user on an empty instance is always allowed and becomes the
    administrator, which closes registration behind them; after that it takes an
    admin reopening it. The settings row is locked across the whole decision so
    two simultaneous first signups cannot both see an empty table and both claim
    administrator.
    """
    try:
        handle = full_handle(localpart)
    except InvalidHandle as exc:
        raise RegistrationError(exc.code, exc.message) from exc

    settings = await app_settings(db, for_update=True)
    bootstrap = not await _has_users(db)
    if not bootstrap and not settings.registration_open:
        raise RegistrationError("registration_closed", "registration is closed on this server")

    user = User(
        handle=handle,
        password_hash=hash_password(password),
        recovery_email=(recovery_email.strip().lower() if recovery_email else None),
        is_admin=bootstrap,
        # Snapshotted at creation: later changes to the instance default leave
        # this user's quota alone.
        quota_bytes=settings.default_quota_bytes,
    )
    db.add(user)
    try:
        await db.flush()
    except IntegrityError as exc:
        await db.rollback()
        raise RegistrationError("handle_taken", "that handle is already registered") from exc
    if bootstrap:
        settings.registration_open = False
    return user


async def _bump_rate_limit(key: str) -> int:
    redis = get_redis()
    count = await redis.incr(key)
    if count == 1:
        await redis.expire(key, _RL_WINDOW)
    return int(count)


async def _check_rate_limit(handle: str, ip: str) -> None:
    for key in (f"rl:login:handle:{handle}", f"rl:login:ip:{ip}"):
        count = await _bump_rate_limit(key)
        if count > _RL_MAX:
            ttl = await get_redis().ttl(key)
            raise RateLimited(retry_after=max(ttl, 1))


async def authenticate(
    db: AsyncSession,
    *,
    handle_input: str,
    password: str,
    ip: str,
    totp: str | None = None,
    trusted_device: Callable[[uuid.UUID], Awaitable[bool]] | None = None,
) -> User | None:
    """Verify credentials (constant-time). Returns the user or None; may rate-limit.

    When the account has TOTP enabled, raises :class:`TotpRequired` (no code given)
    or :class:`TotpInvalid` (wrong code) after the password check passes — unless
    ``trusted_device`` says this browser was trusted for the account at an earlier
    2FA sign-in, in which case the password is enough.
    """
    handle = resolve_handle(handle_input)
    await _check_rate_limit(handle, ip)

    user = await db.scalar(select(User).where(User.handle == handle))
    # Always run a verify to keep timing uniform whether or not the user exists.
    stored_hash = user.password_hash if user else _DUMMY_HASH
    ok = verify_password(stored_hash, password)
    if user is None or not ok:
        return None
    # Only after the password: a stranger must not learn that a handle is blocked.
    if user.blocked_at is not None:
        raise AccountBlocked

    if user.totp_secret is not None:
        from app.identity.totp import decrypt_totp_secret
        from app.identity.totp import verify as verify_totp

        if not totp and trusted_device is not None and await trusted_device(user.id):
            return user
        if not totp:
            raise TotpRequired
        if not verify_totp(decrypt_totp_secret(user.totp_secret), totp):
            raise TotpInvalid
    return user


# A precomputed hash so authentication of a non-existent handle still spends
# argon2 time (mitigates user-enumeration via timing, SPEC §6.1 constant-time).
_DUMMY_HASH = hash_password("x" * 24)