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

from typing import Annotated

from fastapi import APIRouter, 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.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.passwords import hash_password, verify_password
from app.identity.reset import consume_token, issue_token
from app.identity.schemas import (
    LoginRequest,
    MessageOut,
    PasswordResetIn,
    PasswordResetRequestIn,
    RegisterRequest,
    SessionOut,
    TotpDisableRequest,
    TotpEnableRequest,
    TotpSetupResponse,
    UserOut,
)
from app.identity.service import (
    RateLimited,
    RegistrationError,
    TotpInvalid,
    TotpRequired,
    authenticate,
    register_user,
    resolve_handle,
)
from app.identity.sessions import (
    create_session,
    destroy_all_sessions,
    destroy_session,
    list_sessions,
    revoke_session,
)
from app.models import User

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

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,
    )


@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:
        status = 409 if exc.code == "handle_taken" else 400
        raise_api_error(status, 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")
    if user is None:
        raise_api_error(401, "invalid_credentials", "incorrect handle or password")
    await _establish_session(request, response, user)
    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) ---------------------------------------------


@router.post("/password/reset-request", response_model=MessageOut)
async def password_reset_request(body: PasswordResetRequestIn, db: DbSession) -> MessageOut:
    # Always return ok (do not leak whether the handle exists / has recovery mail).
    handle = resolve_handle(body.handle)
    user = await db.scalar(select(User).where(User.handle == handle))
    if user is not None and user.recovery_email:
        await issue_token(user.id, user.recovery_email)
    return MessageOut()


@router.post("/password/reset", response_model=MessageOut)
async def password_reset(body: PasswordResetIn, db: DbSession) -> MessageOut:
    user_id = await consume_token(body.token)
    if user_id is None:
        raise_api_error(400, "invalid_token", "reset token is invalid or expired")
    user = await db.get(User, user_id)
    if user is None:
        raise_api_error(400, "invalid_token", "reset token is invalid or 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()