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
import uuid
from typing import Annotated, NoReturn

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

from app import hosted
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.hosted.admin import MailServerError
from app.identity import app_passwords, invitations, trusted_devices
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.password_policy import PasswordRejected, check_new_password
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 (
    AppPasswordCreatedOut,
    AppPasswordIn,
    AppPasswordOut,
    InvitationCheckIn,
    InvitationCheckOut,
    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,
        mail_host=hosted.mail_hostname() if hosted.enabled() else None,
    )


@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, domain=get_settings().app_domain
    )


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


@router.post("/invitation/check", response_model=InvitationCheckOut)
async def invitation_check(body: InvitationCheckIn, db: DbSession) -> InvitationCheckOut:
    """Whether an invitation link still works, and who it is from — for the
    sign-up page. Does not use it up."""
    if not (await app_settings(db)).invitations_enabled:
        return InvitationCheckOut(valid=False)
    invitation = await invitations.find_valid(db, body.token)
    if invitation is None:
        return InvitationCheckOut(valid=False)
    inviter = await db.get(User, invitation.inviter_id)
    return InvitationCheckOut(
        valid=True,
        invitee_name=invitation.invitee_name,
        inviter_name=await invitations.inviter_name(db, inviter) if inviter else "",
    )


@router.post("/register", response_model=UserOut)
async def register(
    body: RegisterRequest, request: Request, response: Response, db: DbSession
) -> UserOut:
    try:
        await check_new_password(body.password, resolve_handle(body.localpart))
    except PasswordRejected as exc:
        raise_api_error(400, exc.code, exc.message)
    invitation_id: uuid.UUID | None = None
    if body.invite:
        # Used up in this same transaction: if the sign-up fails (the address
        # is taken, …) the rollback hands the invitation back.
        if (await app_settings(db)).invitations_enabled:
            invitation_id = await invitations.claim(db, body.invite)
        if invitation_id is None:
            raise_api_error(
                400, "invalid_invitation", "this invitation has already been used or has expired"
            )
    try:
        user = await register_user(
            db,
            localpart=body.localpart,
            password=body.password,
            recovery_email=body.recovery_email,
            invited=invitation_id is not None,
        )
    except RegistrationError as exc:
        raise_api_error(_REGISTRATION_ERROR_STATUS.get(exc.code, 400), exc.code, exc.message)
    if invitation_id is not None:
        await invitations.record_user(db, invitation_id, user.id)
    await db.commit()
    # The user's own mailbox (App. D §D.4): it exists on the mail server as soon
    # as the user does (the directory reads the users table); this adds it to
    # the webmail.
    from app.hosted import native

    await native.ensure_and_sync(db, user)
    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:
    async def trusted(user_id: uuid.UUID) -> bool:
        return await trusted_devices.is_trusted(request, response, user_id)

    try:
        user = await authenticate(
            db,
            handle_input=body.handle,
            password=body.password,
            ip=_client_ip(request),
            totp=body.totp,
            trusted_device=trusted,
        )
    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")
    if body.remember_device and body.totp and user.totp_secret is not None:
        # Only after a code was actually checked: a trusted sign-in cannot extend
        # the trust to anything new.
        await trusted_devices.trust(request, response, user.id)
    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()
    await trusted_devices.revoke_all(user.id)  # trust was in place of a code that no longer exists
    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")
    # Before the 2FA code and the link are spent: a refused password costs neither.
    try:
        await check_new_password(body.password, user.handle)
    except PasswordRejected as exc:
        raise_api_error(400, exc.code, exc.message)
    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
    # Mail apps sign in again too: the reset may be the owner taking the account back.
    try:
        await app_passwords.revoke_all(db, user.id)
    except MailServerError:
        # The new password is set; the app passwords stay listed (and working)
        # until the mail server answers, so the user can still revoke them.
        logger.exception("password reset: could not revoke app passwords on the mail server")
    # Whoever reset the password may be the owner taking the account back: no
    # device keeps skipping the code on the old trust.
    await trusted_devices.revoke_all(user.id)
    return MessageOut()


# --- App passwords (SPEC App. D §D.5) -----------------------------------------


def _mail_server_down() -> NoReturn:
    raise_api_error(
        502, "mail_server_unavailable", "the mail server is not answering; try again shortly"
    )


def _app_password_out(row: object) -> AppPasswordOut:
    from app.models import UserAppPassword

    assert isinstance(row, UserAppPassword)
    return AppPasswordOut(
        id=str(row.id),
        name=row.name,
        created_at=row.created_at.isoformat(),
        last_used_at=row.last_used_at.isoformat() if row.last_used_at else None,
    )


@router.get("/app-passwords", response_model=list[AppPasswordOut])
async def list_app_passwords(user: CurrentUser, db: DbSession) -> list[AppPasswordOut]:
    return [_app_password_out(r) for r in await app_passwords.list_for(db, user.id)]


@router.post(
    "/app-passwords", response_model=AppPasswordCreatedOut, dependencies=[Depends(require_csrf)]
)
async def create_app_password(
    body: AppPasswordIn, user: CurrentUser, db: DbSession
) -> AppPasswordCreatedOut:
    """A new password for one mail app. Takes the account password: an app
    password opens the mailbox without a 2FA code."""
    if not verify_password(user.password_hash, body.password):
        raise_api_error(403, "invalid_password", "incorrect password")
    if not body.name.strip():
        raise_api_error(400, "invalid_name", "give the app password a name")
    try:
        row, password = await app_passwords.create(db, user.id, body.name)
    except app_passwords.TooMany:
        raise_api_error(
            409, "too_many", f"you can have up to {app_passwords.MAX_PER_USER} app passwords"
        )
    except MailServerError:
        _mail_server_down()
    return AppPasswordCreatedOut(**_app_password_out(row).model_dump(), app_password=password)


@router.delete(
    "/app-passwords/{app_password_id}",
    response_model=MessageOut,
    dependencies=[Depends(require_csrf)],
)
async def revoke_app_password(
    app_password_id: uuid.UUID, user: CurrentUser, db: DbSession
) -> MessageOut:
    try:
        revoked = await app_passwords.revoke(db, user.id, app_password_id)
    except MailServerError:
        _mail_server_down()
    if not revoked:
        raise_api_error(404, "not_found", "app password not found")
    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)