LibraryAPI/library_service/routers/auth.py

"""Модуль работы с авторизацией и аутентификацией пользователей"""

import base64

from datetime import timedelta
from typing import Annotated

from fastapi import APIRouter, Body, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordRequestForm
from sqlmodel import Session, select

from library_service.services import require_captcha
from library_service.models.db import Role, User
from library_service.models.dto import (
    Token,
    UserCreate,
    UserRead,
    UserUpdate,
    UserList,
    RoleRead,
    RoleList,
    Token,
    PartialToken,
    LoginResponse,
    RecoveryCodeUse,
    RegisterResponse,
    RecoveryCodesStatus,
    RecoveryCodesResponse,
    PasswordResetResponse,
    TOTPSetupResponse,
    TOTPVerifyRequest,
    TOTPDisableRequest,
)

from library_service.settings import get_session
from library_service.auth import (
    ACCESS_TOKEN_EXPIRE_MINUTES,
    RequireAuth,
    RequireAdmin,
    RequireStaff,
    authenticate_user,
    get_password_hash,
    decode_token,
    create_access_token,
    create_refresh_token,
    generate_totp_setup,
    generate_codes_for_user,
    verify_and_use_code,
    get_codes_status,
    verify_totp_code,
    verify_password,
    qr_to_bitmap_b64,
    create_partial_token,
    RequirePartialAuth,
    verify_and_use_code,
    cipher,
)


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


@router.post(
    "/register",
    response_model=RegisterResponse,
    status_code=status.HTTP_201_CREATED,
    summary="Регистрация нового пользователя",
    description="Создает нового пользователя и возвращает резервные коды",
)
def register(
    user_data: UserCreate,
    _=Depends(require_captcha),
    session: Session = Depends(get_session),
):
    """Регистрирует нового пользователя в системе"""
    existing_user = session.exec(
        select(User).where(User.username == user_data.username)
    ).first()
    if existing_user:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Username already registered",
        )

    existing_email = session.exec(
        select(User).where(User.email == user_data.email)
    ).first()
    if existing_email:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Email already registered",
        )

    db_user = User(
        **user_data.model_dump(exclude={"password"}),
        hashed_password=get_password_hash(user_data.password),
    )

    default_role = session.exec(select(Role).where(Role.name == "member")).first()
    if default_role:
        db_user.roles.append(default_role)

    session.add(db_user)
    session.commit()
    session.refresh(db_user)

    recovery_codes = generate_codes_for_user(session, db_user)

    return RegisterResponse(
        user=UserRead(
            **db_user.model_dump(),
            roles=[role.name for role in db_user.roles],
        ),
        recovery_codes=RecoveryCodesResponse(
            codes=recovery_codes,
            generated_at=db_user.recovery_codes_generated_at,
        ),
    )


@router.post(
    "/token",
    response_model=LoginResponse,
    summary="Получение токена",
    description="Аутентификация и получение токенов",
)
def login(
    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
    session: Session = Depends(get_session),
):
    """Аутентифицирует пользователя и возвращает JWT токены"""
    user = authenticate_user(session, form_data.username, form_data.password)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Bearer"},
        )

    token_data = {"sub": user.username, "user_id": user.id}

    if user.is_2fa_enabled:
        return LoginResponse(
            partial_token=create_partial_token(token_data),
            token_type="partial",
            requires_2fa=True,
        )

    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    return LoginResponse(
        access_token=create_access_token(
            data=token_data, expires_delta=access_token_expires
        ),
        refresh_token=create_refresh_token(data=token_data),
        token_type="bearer",
        requires_2fa=False,
    )


@router.post(
    "/refresh",
    response_model=Token,
    summary="Обновление токена",
    description="Получение новой пары токенов, используя действующий Refresh токен",
)
def refresh_token(
    refresh_token: str = Body(..., embed=True),
    session: Session = Depends(get_session),
):
    """Обновляет пару токенов (access и refresh)"""
    try:
        token_data = decode_token(refresh_token, expected_type="refresh")
    except HTTPException:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid refresh token",
            headers={"WWW-Authenticate": "Bearer"},
        )

    user = session.get(User, token_data.user_id)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="User not found",
        )

    if not user.is_active:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="User is inactive",
        )

    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    new_access_token = create_access_token(
        data={"sub": user.username, "user_id": user.id},
        expires_delta=access_token_expires,
    )
    new_refresh_token = create_refresh_token(
        data={"sub": user.username, "user_id": user.id}
    )

    return Token(
        access_token=new_access_token,
        refresh_token=new_refresh_token,
        token_type="bearer",
    )


@router.get(
    "/me",
    response_model=UserRead,
    summary="Текущий пользователь",
    description="Получить информацию о текущем авторизованном пользователе",
)
def get_my_profile(current_user: RequireAuth):
    """Возвращает информацию о текущем пользователе"""
    return UserRead(
        **current_user.model_dump(), roles=[role.name for role in current_user.roles]
    )


@router.put(
    "/me",
    response_model=UserRead,
    summary="Обновить профиль",
    description="Обновить информацию текущего пользователя",
)
def update_user_me(
    user_update: UserUpdate,
    current_user: RequireAuth,
    session: Session = Depends(get_session),
):
    """Обновляет профиль текущего пользователя"""
    if user_update.email:
        current_user.email = user_update.email
    if user_update.full_name:
        current_user.full_name = user_update.full_name
    if user_update.password:
        current_user.hashed_password = get_password_hash(user_update.password)

    session.add(current_user)
    session.commit()
    session.refresh(current_user)

    return UserRead(
        **current_user.model_dump(), roles=[role.name for role in current_user.roles]
    )


@router.get(
    "/2fa",
    response_model=TOTPSetupResponse,
    summary="Создание QR-кода TOTP 2FA",
    description="Генерирует секрет и QR-код для настройки TOTP",
)
def get_totp_qr_bitmap(
    current_user: RequireAuth,
    session: Session = Depends(get_session),
):
    """Возвращает данные для настройки TOTP"""
    totp_data = generate_totp_setup(current_user.username)
    encrypted = cipher.encrypt(totp_data["secret"].encode())

    current_user.totp_secret = base64.b64encode(encrypted).decode()
    session.add(current_user)
    session.commit()

    return TOTPSetupResponse(**totp_data)


@router.post(
    "/2fa/enable",
    summary="Включение TOTP 2FA",
    description="Подтверждает настройку и включает 2FA",
)
def enable_2fa(
    data: TOTPVerifyRequest,
    current_user: RequireAuth,
    secret: str = Body(..., embed=True),
    session: Session = Depends(get_session),
):
    """Включает 2FA после проверки кода"""
    if current_user.is_2fa_enabled:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="2FA already enabled",
        )

    if not current_user.totp_secret:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail="Secret key not generated"
        )

    if not verify_totp_code(secret, data.code):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Invalid TOTP code",
        )

    decrypted = cipher.decrypt(base64.b64decode(current_user.totp_secret.encode()))
    if secret != decrypted.decode():
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST, detail="Incorret secret"
        )

    current_user.is_2fa_enabled = True
    session.add(current_user)
    session.commit()

    return {"success": True}


@router.post(
    "/2fa/disable",
    summary="Отключение TOTP 2FA",
    description="Отключает 2FA после проверки пароля и кода",
)
def disable_2fa(
    data: TOTPDisableRequest,
    current_user: RequireAuth,
    session: Session = Depends(get_session),
):
    """Отключает 2FA"""
    if not current_user.is_2fa_enabled or not current_user.totp_secret:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="2FA not enabled",
        )

    if not verify_password(data.password, current_user.hashed_password):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid password",
        )

    current_user.totp_secret = None
    current_user.is_2fa_enabled = False
    session.add(current_user)
    session.commit()

    return {"success": True}


@router.post(
    "/2fa/verify",
    response_model=Token,
    summary="Верификация 2FA",
    description="Завершает аутентификацию с помощью TOTP кода или резервного кода",
)
def verify_2fa(
    data: TOTPVerifyRequest,
    user: RequirePartialAuth,
    session: Session = Depends(get_session),
):
    """Верифицирует 2FA и возвращает полный токен"""
    if not data.code:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Provide TOTP code",
        )

    verified = False

    if data.code and user.totp_secret:
        decrypted = cipher.decrypt(base64.b64decode(user.totp_secret.encode()))
        if verify_totp_code(decrypted.decode(), data.code):
            verified = True

    if not verified:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid 2FA code",
        )

    token_data = {"sub": user.username, "user_id": user.id}
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)

    return Token(
        access_token=create_access_token(
            data=token_data, expires_delta=access_token_expires
        ),
        refresh_token=create_refresh_token(data=token_data),
    )


@router.get(
    "/recovery-codes/status",
    response_model=RecoveryCodesStatus,
    summary="Статус резервных кодов",
    description="Показывает количество оставшихся кодов и какие использованы",
)
def get_recovery_codes_status(current_user: RequireAuth):
    """Возвращает статус резервных кодов"""
    return RecoveryCodesStatus(**get_codes_status(current_user))


@router.post(
    "/recovery-codes/regenerate",
    response_model=RecoveryCodesResponse,
    summary="Перегенерация резервных кодов",
    description="Генерирует новые коды, старые аннулируются",
)
def regenerate_recovery_codes(
    current_user: RequireAuth,
    session: Session = Depends(get_session),
):
    """Генерирует новые резервные коды"""
    codes = generate_codes_for_user(session, current_user)

    return RecoveryCodesResponse(
        codes=codes,
        generated_at=current_user.recovery_codes_generated_at,
    )


@router.post(
    "/password/reset",
    response_model=PasswordResetResponse,
    summary="Сброс пароля через резервный код",
    description="Устанавливает новый пароль используя резервный код",
)
def reset_password(
    data: RecoveryCodeUse,
    session: Session = Depends(get_session),
):
    """Сброс пароля с использованием резервного кода"""
    user = session.exec(select(User).where(User.username == data.username)).first()

    if not user:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Invalid username or recovery code",
        )

    if not user.is_active:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Account is deactivated",
        )

    if not verify_and_use_code(session, user, data.recovery_code):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Invalid username or recovery code",
        )

    user.hashed_password = get_password_hash(data.new_password)
    session.add(user)
    session.commit()

    return PasswordResetResponse(**get_codes_status(user))