stock/backend/app/auth_api.py

"""只登录、不注册的鉴权 API。"""
from __future__ import annotations

from datetime import timedelta

from fastapi import APIRouter, Cookie, Depends, HTTPException, Request, Response, status
from sqlalchemy import delete, select, update
from sqlalchemy.ext.asyncio import AsyncSession

from .auth import (
    drop_session_cache,
    get_auth_session,
    new_session_token,
    require_user,
    session_expiry,
    token_digest,
    utcnow,
    verify_dummy_password,
    verify_password,
)
from .config import settings
from .db import get_session
from .models import AuthSession, User
from .schemas import CurrentUserOut, LoginRequest, LoginResponse


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


def set_session_cookie(response: Response, token: str) -> None:
    response.set_cookie(
        key=settings.auth_cookie_name,
        value=token,
        max_age=settings.auth_session_hours * 60 * 60,
        path="/api",
        secure=settings.auth_cookie_secure,
        httponly=True,
        samesite="strict",
    )


def clear_session_cookie(response: Response) -> None:
    response.delete_cookie(
        key=settings.auth_cookie_name,
        path="/api",
        secure=settings.auth_cookie_secure,
        httponly=True,
        samesite="strict",
    )


@router.post("/login", response_model=LoginResponse)
async def login(
    payload: LoginRequest,
    request: Request,
    response: Response,
    db: AsyncSession = Depends(get_session),
) -> LoginResponse:
    now = utcnow()
    username = payload.username.strip()
    user = (await db.execute(select(User).where(User.username == username))).scalar_one_or_none()

    if user is None:
        verify_dummy_password(payload.password)
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="用户名或密码错误")

    if not user.is_active:
        verify_dummy_password(payload.password)
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="用户名或密码错误")

    if user.locked_until is not None and user.locked_until > now:
        raise HTTPException(
            status_code=status.HTTP_423_LOCKED,
            detail=f"登录失败次数过多，请在 {user.locked_until.isoformat()} 后重试",
        )

    if not verify_password(user.password_hash, payload.password):
        user.failed_login_count += 1
        if user.failed_login_count >= settings.auth_max_failed_logins:
            user.failed_login_count = 0
            user.locked_until = now + timedelta(minutes=settings.auth_lock_minutes)
        user.updated_at = now
        await db.commit()
        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="用户名或密码错误")

    user.failed_login_count = 0
    user.locked_until = None
    user.last_login_at = now
    user.updated_at = now

    # 登录时顺便清理过期或已吊销会话，避免单用户长期运行累积垃圾数据。
    await db.execute(
        delete(AuthSession).where(
            (AuthSession.expires_at <= now) | (AuthSession.revoked_at.is_not(None))
        )
    )

    token = new_session_token()
    expires_at = session_expiry()
    db.add(
        AuthSession(
            user_id=user.id,
            token_hash=token_digest(token),
            created_at=now,
            expires_at=expires_at,
            last_seen_at=now,
            ip_address=request.client.host if request.client else None,
            user_agent=request.headers.get("user-agent", "")[:512] or None,
        )
    )
    await db.commit()
    set_session_cookie(response, token)
    response.headers["Cache-Control"] = "no-store"
    return LoginResponse(user=CurrentUserOut.model_validate(user), expires_at=expires_at)


@router.get("/me", response_model=CurrentUserOut)
async def me(user: User = Depends(require_user)) -> CurrentUserOut:
    return CurrentUserOut.model_validate(user)


@router.post("/logout", status_code=status.HTTP_204_NO_CONTENT)
async def logout(
    response: Response,
    stock_session: str | None = Cookie(default=None, alias=settings.auth_cookie_name),
    db: AsyncSession = Depends(get_session),
) -> Response:
    if stock_session:
        await db.execute(
            update(AuthSession)
            .where(AuthSession.token_hash == token_digest(stock_session))
            .values(revoked_at=utcnow())
        )
        await db.commit()
        drop_session_cache(digest=token_digest(stock_session))  # 同进程立即失效，登出即时生效
    clear_session_cookie(response)
    response.status_code = status.HTTP_204_NO_CONTENT
    response.headers["Cache-Control"] = "no-store"
    return response


@router.post("/logout-all", status_code=status.HTTP_204_NO_CONTENT)
async def logout_all(
    response: Response,
    user: User = Depends(require_user),
    db: AsyncSession = Depends(get_session),
) -> Response:
    await db.execute(
        update(AuthSession)
        .where(AuthSession.user_id == user.id, AuthSession.revoked_at.is_(None))
        .values(revoked_at=utcnow())
    )
    await db.commit()
    drop_session_cache(user_id=user.id)  # 该用户全部 token 的本地缓存立即失效
    clear_session_cookie(response)
    response.status_code = status.HTTP_204_NO_CONTENT
    response.headers["Cache-Control"] = "no-store"
    return response