Улучшение безопасности

2026-02-04 04:31:09 +00:00 · 2026-01-19 23:22:29 +03:00
parent 758e0fc9e6
commit d6ecd4066f
59 changed files with 2712 additions and 1010 deletions
@@ -2,9 +2,11 @@

 from datetime import timedelta
 from typing import Annotated
+from pathlib import Path

 from fastapi import APIRouter, Body, Depends, HTTPException, status, Request
 from fastapi.security import OAuth2PasswordRequestForm
+from fastapi.templating import Jinja2Templates
 from sqlmodel import Session, select
 import pyotp

@@ -17,7 +19,19 @@ from library_service.models.dto import (
    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,
@@ -29,10 +43,18 @@ from library_service.auth import (
    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,
 )
-from pathlib import Path
-from fastapi.templating import Jinja2Templates
+

 templates = Jinja2Templates(directory=Path(__file__).parent.parent / "templates")
 router = APIRouter(prefix="/auth", tags=["authentication"])
@@ -40,10 +62,10 @@ router = APIRouter(prefix="/auth", tags=["authentication"])

@router.post(
    "/register",
-    response_model=UserRead,
+    response_model=RegisterResponse,
    status_code=status.HTTP_201_CREATED,
    summary="Регистрация нового пользователя",
-    description="Создает нового пользователя в системе",
+    description="Создает нового пользователя и возвращает резервные коды",
 )
 def register(user_data: UserCreate, session: Session = Depends(get_session)):
    """Регистрирует нового пользователя в системе"""
@@ -61,7 +83,8 @@ def register(user_data: UserCreate, session: Session = Depends(get_session)):
    ).first()
    if existing_email:
        raise HTTPException(
-            status_code=status.HTTP_400_BAD_REQUEST, detail="Email already registered"
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Email already registered",
        )

    db_user = User(
@@ -77,14 +100,25 @@ def register(user_data: UserCreate, session: Session = Depends(get_session)):
    session.commit()
    session.refresh(db_user)

-    return UserRead(**db_user.model_dump(), roles=[role.name for role in db_user.roles])
+    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=Token,
+    response_model=LoginResponse,
    summary="Получение токена",
-    description="Аутентификация и получение JWT токена",
+    description="Аутентификация и получение токенов",
 )
 def login(
    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
@@ -99,17 +133,23 @@ def login(
            headers={"WWW-Authenticate": "Bearer"},
        )

-    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
-    access_token = create_access_token(
-        data={"sub": user.username, "user_id": user.id},
-        expires_delta=access_token_expires,
-    )
-    refresh_token = create_refresh_token(
-        data={"sub": user.username, "user_id": user.id}
-    )
+    token_data = {"sub": user.username, "user_id": user.id}

-    return Token(
-        access_token=access_token, refresh_token=refresh_token, token_type="bearer"
+    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,
    )


@@ -330,18 +370,182 @@ def get_roles(

@router.get(
    "/2fa",
+    response_model=TOTPSetupResponse,
    summary="Создание QR-кода TOTP 2FA",
-    description="Получить информацию о текущем авторизованном пользователе",
+    description="Генерирует секрет и QR-код для настройки TOTP",
 )
 def get_totp_qr_bitmap(auth: RequireAuth):
-    """Возвращает qr-код bitmap"""
-    issuer = "issuer"
-    username = auth.username
-    secret = pyotp.random_base32()
+    """Возвращает данные для настройки TOTP"""
+    return TOTPSetupResponse(**generate_totp_setup(auth.username))

-    totp = pyotp.TOTP(secret)
-    provisioning_uri = totp.provisioning_uri(name=username, issuer_name=issuer)

-    bitmap_data = qr_to_bitmap_b64(provisioning_uri)
+@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",
+        )

-    return {"secret": secret, "username": username, "issuer": issuer, **bitmap_data}
+    if not verify_totp_code(secret, data.code):
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Invalid TOTP code",
+        )
+
+    current_user.totp_secret = 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:
+        if verify_totp_code(user.totp_secret, 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))