Budget_Tool_Martinhal/backend/app/routers/auth.py

from datetime import datetime, timezone
from io import BytesIO

from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, UploadFile
from pydantic import BaseModel
from sqlalchemy import func, select
from sqlalchemy.orm import Session

from .. import mfa
from ..audit import record
from ..config import settings
from ..db import get_db
from ..models import User
from ..security import COOKIE, current_user, hash_password, make_token, verify_password
from ..serializers import user_out

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


class LoginIn(BaseModel):
    email: str
    password: str
    code: str | None = None  # authenticator or recovery code, for users with MFA on


class ProfileIn(BaseModel):
    name: str | None = None
    language: str | None = None
    theme: str | None = None
    current_password: str | None = None
    new_password: str | None = None


@router.post("/login")
def login(body: LoginIn, response: Response, db: Session = Depends(get_db)):
    user = db.scalar(select(User).where(func.lower(User.email) == body.email.strip().lower()))
    if not user or not user.active or not verify_password(body.password, user.password_hash):
        record(db, user if user and user.active else None, "user", user.id if user else None, "login_failed",
               body.email.strip()[:200])
        db.commit()
        raise HTTPException(401, "invalid_credentials")
    if user.mfa_enabled:
        if not (body.code or "").strip():
            # password is right: the sign-in screen now asks for the code
            return {"mfa_required": True}
        if mfa.locked(user):
            raise HTTPException(429, "mfa_locked")
        kind = mfa.check(user, body.code)
        if not kind:
            record(db, user, "user", user.id, "mfa_failed", user.email)
            db.commit()
            raise HTTPException(429 if mfa.locked(user) else 401, "mfa_locked" if mfa.locked(user) else "invalid_mfa_code")
        if kind == "recovery":
            record(db, user, "user", user.id, "mfa_recovery_used", f"{len(user.mfa_recovery or [])} recovery codes left")
    response.set_cookie(COOKIE, make_token(user), httponly=True, samesite="lax",
                        secure=settings.cookie_secure, max_age=settings.session_hours * 3600)
    record(db, user, "user", user.id, "login")
    db.commit()
    return user_out(user)


@router.post("/logout")
def logout(request: Request, response: Response, db: Session = Depends(get_db)):
    try:
        user = current_user(request, db)
        record(db, user, "user", user.id, "logout")
        db.commit()
    except HTTPException:
        pass
    response.delete_cookie(COOKIE)
    return {"ok": True}


@router.get("/me")
def me(user: User = Depends(current_user)):
    return user_out(user)


@router.put("/me")
def update_me(body: ProfileIn, user: User = Depends(current_user), db: Session = Depends(get_db)):
    if body.name:
        user.name = body.name.strip()
    if body.language in ("pt", "en"):
        user.language = body.language
    if body.theme in ("system", "light", "dark"):
        user.theme = body.theme
    if body.new_password:
        if not body.current_password or not verify_password(body.current_password, user.password_hash):
            raise HTTPException(400, "wrong_password")
        if len(body.new_password) < 8:
            raise HTTPException(400, "password_too_short")
        user.password_hash = hash_password(body.new_password)
        record(db, user, "user", user.id, "password_changed")
    db.commit()
    return user_out(user)


# --- MFA ----------------------------------------------------------------------------------------

class CodeIn(BaseModel):
    code: str = ""
    password: str = ""


@router.post("/mfa/setup")
def mfa_setup(user: User = Depends(current_user), db: Session = Depends(get_db)):
    """Starts MFA set-up: a new secret to scan as a QR code. MFA is only on after a code is confirmed."""
    if user.mfa_enabled:
        raise HTTPException(409, "mfa_already_on")
    secret = mfa.new_secret()
    user.mfa_pending_secret = mfa.seal(secret)
    db.commit()
    uri = mfa.otpauth_uri(secret, user.email)
    return {"secret": secret, "uri": uri, "qr_svg": mfa.qr_svg(uri)}


@router.post("/mfa/enable")
def mfa_enable(body: CodeIn, user: User = Depends(current_user), db: Session = Depends(get_db)):
    secret = mfa.unseal(user.mfa_pending_secret)
    if user.mfa_enabled or not secret:
        raise HTTPException(409, "mfa_setup_first")
    step = mfa.matching_step(secret, body.code)
    if step is None:
        raise HTTPException(400, "invalid_mfa_code")
    codes, hashes = mfa.new_recovery_codes()
    user.mfa_secret, user.mfa_pending_secret, user.mfa_enabled = user.mfa_pending_secret, None, True
    user.mfa_enabled_at, user.mfa_last_step, user.mfa_recovery = datetime.now(timezone.utc), step, hashes
    user.mfa_failed, user.mfa_locked_until = 0, None
    record(db, user, "user", user.id, "mfa_enabled", user.email)
    db.commit()
    return {"user": user_out(user), "recovery_codes": codes}


@router.post("/mfa/recovery-codes")
def mfa_new_recovery_codes(body: CodeIn, user: User = Depends(current_user), db: Session = Depends(get_db)):
    """New recovery codes (the old ones stop working). Needs a current code."""
    if not user.mfa_enabled:
        raise HTTPException(409, "mfa_off")
    if not mfa.check(user, body.code):
        db.commit()
        raise HTTPException(400, "invalid_mfa_code")
    codes, user.mfa_recovery = mfa.new_recovery_codes()
    record(db, user, "user", user.id, "mfa_recovery_renewed", user.email)
    db.commit()
    return {"user": user_out(user), "recovery_codes": codes}


@router.post("/mfa/disable")
def mfa_disable(body: CodeIn, user: User = Depends(current_user), db: Session = Depends(get_db)):
    """Turning MFA off needs the password and a current code (or a recovery code)."""
    if not user.mfa_enabled:
        raise HTTPException(409, "mfa_off")
    if not verify_password(body.password, user.password_hash):
        raise HTTPException(400, "wrong_password")
    if not mfa.check(user, body.code):
        db.commit()
        raise HTTPException(400, "invalid_mfa_code")
    mfa.disable(user)
    record(db, user, "user", user.id, "mfa_disabled", user.email)
    db.commit()
    return user_out(user)


# --- profile picture ----------------------------------------------------------------------------

AVATAR_MAX_BYTES = 5 * 1024 * 1024
AVATAR_SIZE = 256


def resize_avatar(data: bytes) -> tuple[bytes, str]:
    """Any common picture -> a square JPEG (or PNG when it has transparency) of at most 256 px."""
    from PIL import Image, ImageOps, UnidentifiedImageError
    try:
        img = Image.open(BytesIO(data))
        img.load()
    except (UnidentifiedImageError, OSError, Image.DecompressionBombError):
        raise HTTPException(400, "not_an_image")
    img = ImageOps.exif_transpose(img)
    side = min(img.size)
    left, top = (img.width - side) // 2, (img.height - side) // 2
    img = img.crop((left, top, left + side, top + side))
    img.thumbnail((AVATAR_SIZE, AVATAR_SIZE))
    out = BytesIO()
    if img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info):
        img.convert("RGBA").save(out, "PNG", optimize=True)
        return out.getvalue(), "image/png"
    img.convert("RGB").save(out, "JPEG", quality=88, optimize=True)
    return out.getvalue(), "image/jpeg"


@router.put("/me/avatar")
async def upload_avatar(file: UploadFile = File(...), user: User = Depends(current_user), db: Session = Depends(get_db)):
    data = await file.read(AVATAR_MAX_BYTES + 1)
    if len(data) > AVATAR_MAX_BYTES:
        raise HTTPException(413, "image_too_large")
    user.avatar, user.avatar_type = resize_avatar(data)
    user.avatar_updated_at = datetime.now(timezone.utc)
    record(db, user, "user", user.id, "avatar_changed", user.email)
    db.commit()
    return user_out(user)


@router.delete("/me/avatar")
def delete_avatar(user: User = Depends(current_user), db: Session = Depends(get_db)):
    user.avatar = user.avatar_type = user.avatar_updated_at = None
    record(db, user, "user", user.id, "avatar_removed", user.email)
    db.commit()
    return user_out(user)