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

"""Auth API routes (SPEC §9.1): register, login (2FA-aware), logout, me,
TOTP setup/enable, session list/revoke, password reset."""

from __future__ import annotations

import logging
import re
from typing import Annotated

from fastapi import APIRouter, BackgroundTasks, Depends, Request, Response
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.errors import raise_api_error
from app.config import get_settings
from app.db import get_session
from app.events import get_event_bus
from app.identity import totp as totp_mod
from app.identity.cookies import (
    clear_auth_cookies,
    new_csrf_token,
    set_auth_cookies,
)
from app.identity.deps import CurrentSession, CurrentUser, require_csrf
from app.identity.mailer import MailerConfig, MailerError, mailer_config, reset_mail, send_mail
from app.identity.passwords import hash_password, verify_password
from app.identity.reset import (
    consume_token,
    issue_token,
    peek_token,
    record_totp_failure,
    request_allowed,
)
from app.identity.schemas import (
    LoginRequest,
    MessageOut,
    PasswordResetCheckIn,
    PasswordResetCheckOut,
    PasswordResetIn,
    PasswordResetRequestIn,
    RecoveryEmailIn,
    RegisterRequest,
    RegistrationStatusOut,
    SessionOut,
    TotpDisableRequest,
    TotpEnableRequest,
    TotpSetupResponse,
    UserOut,
)
from app.identity.service import (
    AccountBlocked,
    RateLimited,
    RegistrationError,
    TotpInvalid,
    TotpRequired,
    app_settings,
    authenticate,
    register_user,
    registration_status,
    resolve_handle,
)
from app.identity.sessions import (
    create_session,
    destroy_all_sessions,
    destroy_session,
    list_sessions,
    revoke_session,
)
from app.models import User
from app.redis_client import get_redis
from app.sync import leases

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/auth", tags=["auth"])

# Deliberately loose: the recovery address only has to be something mail can be
# sent to; the mail server is the real judge.
_EMAIL_RE = re.compile(r"[^@\s]+@[^@\s]+\.[^@\s]+")

DbSession = Annotated[AsyncSession, Depends(get_session)]


def _client_ip(request: Request) -> str:
    return request.client.host if request.client else "unknown"


async def _establish_session(request: Request, response: Response, user: User) -> None:
    token = await create_session(user.id, request.headers.get("user-agent", ""))
    csrf = new_csrf_token()
    max_age = get_settings().session_ttl_days * 24 * 3600
    set_auth_cookies(request, response, token, csrf, max_age)


def _user_out(user: User) -> UserOut:
    return UserOut(
        id=user.id,
        handle=user.handle,
        recovery_email=user.recovery_email,
        settings=user.settings,
        totp_enabled=user.totp_secret is not None,
        is_admin=user.is_admin,
    )


@router.get("/registration", response_model=RegistrationStatusOut)
async def registration(db: DbSession) -> RegistrationStatusOut:
    """Whether signup is open. Unauthenticated: the signup page needs it to render."""
    status = await registration_status(db)
    return RegistrationStatusOut(open=status.open, bootstrap=status.bootstrap)


_REGISTRATION_ERROR_STATUS = {"handle_taken": 409, "registration_closed": 403}


@router.post("/register", response_model=UserOut)
async def register(
    body: RegisterRequest, request: Request, response: Response, db: DbSession
) -> UserOut:
    try:
        user = await register_user(
            db,
            localpart=body.localpart,
            password=body.password,
            recovery_email=body.recovery_email,
        )
    except RegistrationError as exc:
        raise_api_error(_REGISTRATION_ERROR_STATUS.get(exc.code, 400), exc.code, exc.message)
    await db.commit()
    await _establish_session(request, response, user)
    return _user_out(user)


@router.post("/login", response_model=UserOut)
async def login(body: LoginRequest, request: Request, response: Response, db: DbSession) -> UserOut:
    try:
        user = await authenticate(
            db,
            handle_input=body.handle,
            password=body.password,
            ip=_client_ip(request),
            totp=body.totp,
        )
    except RateLimited as exc:
        response.headers["Retry-After"] = str(exc.retry_after)
        raise_api_error(429, "rate_limited", "too many attempts; try again later")
    except TotpRequired:
        raise_api_error(401, "totp_required", "a 2FA code is required")
    except TotpInvalid:
        raise_api_error(401, "totp_invalid", "incorrect 2FA code")
    except AccountBlocked:
        raise_api_error(
            403, "account_blocked", "this account has been blocked by the administrator"
        )
    if user is None:
        raise_api_error(401, "invalid_credentials", "incorrect handle or password")
    await _establish_session(request, response, user)
    # Sync v2 §3/§5: a login marks the user active (session heat) and promotes
    # their accounts' tier to hot.
    settings = get_settings()
    await leases.mark_user_active(get_redis(), user.id, settings.tier_session_grace_min * 60)
    await get_event_bus().publish("user.active", {"user_id": str(user.id)})
    return _user_out(user)


@router.post("/logout", response_model=MessageOut, dependencies=[Depends(require_csrf)])
async def logout(request: Request, response: Response, session: CurrentSession) -> MessageOut:
    await destroy_session(session.token)
    clear_auth_cookies(request, response)
    return MessageOut()


@router.get("/me", response_model=UserOut)
async def me(user: CurrentUser) -> UserOut:
    return _user_out(user)


# --- Two-factor authentication (SPEC §6.1) ----------------------------------


@router.post("/totp/setup", response_model=TotpSetupResponse, dependencies=[Depends(require_csrf)])
async def totp_setup(user: CurrentUser) -> TotpSetupResponse:
    if user.totp_secret is not None:
        raise_api_error(409, "totp_already_enabled", "2FA is already enabled")
    secret = await totp_mod.start_setup(user.id)
    return TotpSetupResponse(
        secret=secret, provisioning_uri=totp_mod.provisioning_uri(secret, user.handle)
    )


@router.post("/totp/enable", response_model=MessageOut, dependencies=[Depends(require_csrf)])
async def totp_enable(body: TotpEnableRequest, user: CurrentUser, db: DbSession) -> MessageOut:
    secret = await totp_mod.get_pending(user.id)
    if secret is None:
        raise_api_error(400, "totp_no_setup", "start TOTP setup first")
    if not totp_mod.verify(secret, body.code):
        raise_api_error(400, "totp_invalid", "incorrect 2FA code")
    user.totp_secret = totp_mod.encrypt_totp_secret(secret)
    await db.commit()
    await totp_mod.clear_pending(user.id)
    return MessageOut()


@router.post("/totp/disable", response_model=MessageOut, dependencies=[Depends(require_csrf)])
async def totp_disable(body: TotpDisableRequest, user: CurrentUser, db: DbSession) -> MessageOut:
    # Re-authenticate with the password before turning off 2FA.
    if not verify_password(user.password_hash, body.password):
        raise_api_error(401, "invalid_credentials", "incorrect password")
    user.totp_secret = None
    await db.commit()
    return MessageOut()


# --- Session management (SPEC §9.1) -----------------------------------------


@router.get("/sessions", response_model=list[SessionOut])
async def sessions(user: CurrentUser, session: CurrentSession) -> list[SessionOut]:
    infos = await list_sessions(user.id, session.token)
    return [
        SessionOut(sid=i.sid, created=i.created.isoformat(), ua=i.ua, current=i.current)
        for i in infos
    ]


@router.delete("/sessions/{sid}", response_model=MessageOut, dependencies=[Depends(require_csrf)])
async def revoke(sid: str, user: CurrentUser) -> MessageOut:
    if not await revoke_session(user.id, sid):
        raise_api_error(404, "not_found", "session not found")
    return MessageOut()


# --- Password reset (SPEC §6.1) ---------------------------------------------


async def _send_reset_mail(cfg: MailerConfig, to: str, handle: str, token: str) -> None:
    """Runs after the response has gone out (see password_reset_request)."""
    # The token rides in the fragment, which browsers never send to a server, so
    # it stays out of access logs and Referer headers.
    subject, text, html = reset_mail(f"{cfg.public_url}/reset-password#token={token}", handle)
    try:
        await send_mail(cfg, to, subject, text, html)
    except MailerError as exc:
        logger.warning("password reset mail for %s could not be sent: %s", handle, exc)


@router.post("/password/reset-request", response_model=MessageOut)
async def password_reset_request(
    body: PasswordResetRequestIn, background: BackgroundTasks, db: DbSession
) -> MessageOut:
    """Email a reset link to the account's recovery address.

    The answer is the same whether or not the account exists, has a recovery
    address, or hit the hourly cap, and the mail is sent after the response —
    so neither the reply nor its timing tells a stranger anything.
    """
    handle = resolve_handle(body.handle)
    user = await db.scalar(select(User).where(User.handle == handle))
    if user is None or not user.recovery_email:
        return MessageOut()
    cfg = mailer_config(await app_settings(db))
    if cfg is None:
        logger.warning("password reset requested for %s, but system mail is not set up", handle)
        return MessageOut()
    if not await request_allowed(user.id):
        return MessageOut()
    token = await issue_token(user.id)
    background.add_task(_send_reset_mail, cfg, user.recovery_email, user.handle, token)
    return MessageOut()


@router.post("/password/reset/check", response_model=PasswordResetCheckOut)
async def password_reset_check(body: PasswordResetCheckIn, db: DbSession) -> PasswordResetCheckOut:
    """Whether a reset link still works, so the page can say so before asking
    for a new password. Does not use the link up."""
    user_id = await peek_token(body.token)
    user = await db.get(User, user_id) if user_id else None
    if user is None:
        return PasswordResetCheckOut(valid=False)
    return PasswordResetCheckOut(valid=True, totp_required=user.totp_secret is not None)


@router.post("/password/reset", response_model=MessageOut)
async def password_reset(body: PasswordResetIn, db: DbSession) -> MessageOut:
    """Set a new password with a reset link. With 2FA on, the code is checked
    before anything changes; a missing or wrong code leaves the link usable
    (up to MAX_TOTP_FAILURES wrong codes)."""
    user_id = await peek_token(body.token)
    user = await db.get(User, user_id) if user_id else None
    if user is None:
        raise_api_error(400, "invalid_token", "this link is invalid or has expired")
    if user.totp_secret is not None:
        if not body.totp:
            raise_api_error(401, "totp_required", "a 2FA code is required")
        if not totp_mod.verify(totp_mod.decrypt_totp_secret(user.totp_secret), body.totp):
            await record_totp_failure(body.token)
            raise_api_error(401, "totp_invalid", "incorrect 2FA code")
    # Only now is the link spent — and only if nobody spent it in the meantime.
    if await consume_token(body.token) != user.id:
        raise_api_error(400, "invalid_token", "this link is invalid or has expired")
    user.password_hash = hash_password(body.password)
    await db.commit()
    await destroy_all_sessions(user.id)  # force re-login everywhere after reset
    return MessageOut()


@router.put("/recovery-email", response_model=UserOut, dependencies=[Depends(require_csrf)])
async def set_recovery_email(body: RecoveryEmailIn, user: CurrentUser, db: DbSession) -> UserOut:
    """Add, change or remove the recovery address. Takes the password: whoever
    controls this address can take the account over with a reset link."""
    if not verify_password(user.password_hash, body.password):
        raise_api_error(403, "invalid_password", "incorrect password")
    email = (body.recovery_email or "").strip().lower() or None
    if email is not None and not _EMAIL_RE.fullmatch(email):
        raise_api_error(400, "invalid_email", "enter a valid email address")
    user.recovery_email = email
    await db.commit()
    return _user_out(user)