document-management/backend/dms/services/login_lockout.py

"""登录失败锁定策略：连续密码错误达到阈值后临时锁定账号。

- 仅 /auth/login 的密码错误计入（用户名不存在、账号停用、IP受限不计入）；
- 锁定期间登录直接拒绝（ACCOUNT_LOCKED），不延长锁定窗口；
- 锁定到期自动解除；登录成功即清零计数与锁定。
"""

from __future__ import annotations

from datetime import datetime, timedelta, timezone

from sqlalchemy import select

from dms.database.transaction import for_update, transaction
from dms.extensions import db
from dms.models import User

LOGIN_MAX_FAILURES = 5
LOGIN_LOCK_DURATION = timedelta(minutes=15)


def _utc_now() -> datetime:
    return datetime.now(timezone.utc).replace(tzinfo=None)


def is_locked(user: User) -> bool:
    """账号当前是否处于锁定窗口内（已过期的锁定视为未锁定）。"""
    return user.locked_until is not None and user.locked_until > _utc_now()


def apply_failed_attempt(user: User, *, now: datetime | None = None) -> None:
    """纯字段变换：连续失败计数+1，达到阈值则锁定并清零计数。

    不触碰数据库，方便单元测试；调用方负责 row_version 递增。
    """
    moment = now or _utc_now()
    user.failed_login_count = (user.failed_login_count or 0) + 1
    if user.failed_login_count >= LOGIN_MAX_FAILURES:
        user.locked_until = moment + LOGIN_LOCK_DURATION
        user.failed_login_count = 0


def register_failed_password_attempt(user: User) -> None:
    """密码校验失败后落库计数（独立短事务，对齐 record_login_failure）。"""
    db.session.rollback()
    with transaction() as session:
        row = session.scalar(
            for_update(select(User).where(User.id == user.id))
        )
        if row is None:
            return
        apply_failed_attempt(row)
        row.row_version += 1


def clear_failed_attempts(user: User) -> None:
    """登录成功后清零计数与锁定；由调用方事务负责 row_version 递增。"""
    user.failed_login_count = 0
    user.locked_until = None