Ten31-Portal/backend/ten31portal/routers/auth_router.py

"""Authentication endpoints."""

import os
import secrets
import time
from datetime import datetime

from fastapi import APIRouter, Depends, HTTPException, Request
from sqlmodel import Session, select

from ten31portal import config, totp
from ten31portal.audit import record_audit
from ten31portal.auth import get_current_user, hash_password, verify_password
from ten31portal.database import get_session
from ten31portal.models import User
from ten31portal.ratelimit import SlidingWindowLimiter
from ten31portal.schemas import (
    ChangePasswordRequest,
    LoginPending2FA,
    LoginRequest,
    TotpConfirmRequest,
    TotpConfirmResponse,
    TotpDisableRequest,
    TotpSetupResponse,
    TotpVerifyRequest,
    UserResponse,
)

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

# Throttle password guessing per client IP: at most 10 failed attempts per 5 minutes.
# Second-factor code guesses count against the same window.
_login_limiter = SlidingWindowLimiter(max_attempts=10, window_seconds=300)

# How long a password-accepted session may wait for its second factor.
_PENDING_2FA_MAX_AGE = 300  # seconds

# A throwaway hash verified when no user matches, so a missing account costs the same argon2
# time as a real one — otherwise the timing difference reveals which usernames exist.
_DUMMY_HASH = hash_password(secrets.token_urlsafe(16))


@router.post("/login")
def login(
    body: LoginRequest,
    request: Request,
    session: Session = Depends(get_session),
) -> UserResponse | LoginPending2FA:
    client_ip = request.client.host if request.client else "unknown"
    retry_after = _login_limiter.retry_after(client_ip)
    if retry_after > 0:
        raise HTTPException(
            status_code=429,
            detail="Too many login attempts. Please wait a moment and try again.",
            headers={"Retry-After": str(int(retry_after) + 1)},
        )

    # Accept either a username or an email in the login field.
    handle = body.login.strip()
    user = session.exec(select(User).where(User.username == handle)).first()
    if user is None:
        user = session.exec(select(User).where(User.email == handle)).first()
    # Always run a verify (dummy hash when the user is unknown) so success/failure take the
    # same time, and give one generic message so neither branch leaks whether the user exists.
    if user is None:
        verify_password(body.password, _DUMMY_HASH)
        _login_limiter.record_failure(client_ip)
        raise HTTPException(status_code=401, detail="Invalid username or password")
    if not verify_password(body.password, user.password_hash):
        _login_limiter.record_failure(client_ip)
        raise HTTPException(status_code=401, detail="Invalid username or password")
    # Past this point the password was correct, so these messages don't aid guessing.
    if user.primary_account_id is not None:
        raise HTTPException(
            status_code=401,
            detail="This account signs in under another login. Use that account's credentials.",
        )
    if not user.login_enabled:
        raise HTTPException(status_code=401, detail="This account does not have a login yet.")
    if not user.is_active:
        raise HTTPException(status_code=401, detail="Account disabled")
    # Anyone still signing in with the shared default password gets flagged so the portal
    # forces them to set their own before doing anything else (covers accounts created
    # before the flag existed).
    if body.password == config.DEFAULT_INVESTOR_PASSWORD and not user.must_change_password:
        user.must_change_password = True
        session.add(user)
        session.commit()
    if user.totp_enabled:
        # Password accepted, but don't sign the session in yet — park the login until
        # /login/verify-totp confirms the second factor.
        request.session.clear()
        request.session["pending_2fa_user_id"] = user.id
        request.session["pending_2fa_at"] = time.time()
        return LoginPending2FA()
    _login_limiter.reset(client_ip)
    request.session["user_id"] = user.id
    return UserResponse.model_validate(user, from_attributes=True)


@router.post("/login/verify-totp")
def verify_totp_login(
    body: TotpVerifyRequest,
    request: Request,
    session: Session = Depends(get_session),
) -> UserResponse:
    """Second login step: accept an authenticator code or an unused recovery code."""
    client_ip = request.client.host if request.client else "unknown"
    retry_after = _login_limiter.retry_after(client_ip)
    if retry_after > 0:
        raise HTTPException(
            status_code=429,
            detail="Too many attempts. Please wait a moment and try again.",
            headers={"Retry-After": str(int(retry_after) + 1)},
        )

    user_id = request.session.get("pending_2fa_user_id")
    started = request.session.get("pending_2fa_at", 0)
    if user_id is None or time.time() - started > _PENDING_2FA_MAX_AGE:
        request.session.clear()
        raise HTTPException(status_code=401, detail="Sign-in expired. Please log in again.")
    user = session.get(User, user_id)
    if user is None or not user.is_active or not user.totp_enabled or not user.totp_secret:
        request.session.clear()
        raise HTTPException(status_code=401, detail="Sign-in expired. Please log in again.")

    if not totp.verify_code(user.totp_secret, body.code):
        # Not a current authenticator code — maybe a one-time recovery code.
        remaining = totp.consume_recovery_code(user.totp_recovery_codes, body.code)
        if remaining is None:
            _login_limiter.record_failure(client_ip)
            raise HTTPException(status_code=401, detail="Invalid code. Try again.")
        user.totp_recovery_codes = remaining
        session.add(user)
        record_audit(session, user.id, "totp_recovery_used", "user", user.id)
        session.commit()

    _login_limiter.reset(client_ip)
    request.session.clear()
    request.session["user_id"] = user.id
    return UserResponse.model_validate(user, from_attributes=True)


@router.post("/logout")
def logout(request: Request) -> dict[str, str]:
    request.session.clear()
    return {"status": "ok"}


@router.get("/me")
def me(user: User = Depends(get_current_user)) -> UserResponse:
    return UserResponse.model_validate(user, from_attributes=True)


@router.post("/change-password")
def change_password(
    body: ChangePasswordRequest,
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> dict[str, str]:
    """Let the signed-in user set their own password (after confirming the current one)."""
    if not verify_password(body.current_password, user.password_hash):
        raise HTTPException(status_code=400, detail="Current password is incorrect.")
    if len(body.new_password) < 8:
        raise HTTPException(status_code=400, detail="New password must be at least 8 characters.")
    if body.new_password == config.DEFAULT_INVESTOR_PASSWORD:
        raise HTTPException(status_code=400, detail="Please choose a password of your own.")
    user.password_hash = hash_password(body.new_password)
    user.must_change_password = False
    session.add(user)
    session.commit()
    # Once the built-in admin sets their own password, the initial one from first boot is stale.
    if user.is_service_admin:
        try:
            os.remove(config.ADMIN_PASSWORD_FILE)
        except OSError:
            pass
    return {"status": "ok"}


@router.post("/onboarded")
def mark_onboarded(
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> dict[str, str]:
    """Stamp the first-login welcome flow as finished (or skipped) so it stops showing."""
    if user.onboarded_at is None:
        user.onboarded_at = datetime.utcnow()
        session.add(user)
        session.commit()
    return {"status": "ok"}


# --- Two-factor enrollment (per-user opt-in) ---

@router.post("/totp/setup")
def totp_setup(
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> TotpSetupResponse:
    """Start enrollment: mint a secret and return the QR. Not active until confirmed."""
    if user.totp_enabled:
        raise HTTPException(
            status_code=400, detail="Two-factor is already enabled. Disable it first to re-enroll."
        )
    secret = totp.new_secret()
    user.totp_secret = secret
    session.add(user)
    session.commit()
    uri = totp.otpauth_uri(secret, user.username)
    return TotpSetupResponse(secret=secret, otpauth_uri=uri, qr_svg=totp.qr_svg(uri))


@router.post("/totp/confirm")
def totp_confirm(
    body: TotpConfirmRequest,
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> TotpConfirmResponse:
    """Finish enrollment: prove the authenticator works, then hand out recovery codes."""
    if user.totp_enabled:
        raise HTTPException(status_code=400, detail="Two-factor is already enabled.")
    if not user.totp_secret:
        raise HTTPException(status_code=400, detail="Start two-factor setup first.")
    if not totp.verify_code(user.totp_secret, body.code):
        raise HTTPException(status_code=400, detail="That code didn't match. Try again.")
    codes, digests_json = totp.generate_recovery_codes()
    user.totp_enabled = True
    user.totp_recovery_codes = digests_json
    session.add(user)
    record_audit(session, user.id, "totp_enabled", "user", user.id)
    session.commit()
    return TotpConfirmResponse(recovery_codes=codes)


@router.post("/totp/disable")
def totp_disable(
    body: TotpDisableRequest,
    user: User = Depends(get_current_user),
    session: Session = Depends(get_session),
) -> dict[str, str]:
    """Turn off two-factor (requires the account password, not just a live session)."""
    if not verify_password(body.password, user.password_hash):
        raise HTTPException(status_code=400, detail="Password is incorrect.")
    user.totp_secret = None
    user.totp_enabled = False
    user.totp_recovery_codes = None
    session.add(user)
    record_audit(session, user.id, "totp_disabled", "user", user.id)
    session.commit()
    return {"status": "ok"}