document-management/backend/dms/security/decorators.py

"""DMS Bearer 认证与角色授权装饰器。"""

from __future__ import annotations

from functools import wraps
from typing import Any, Callable, TypeVar, cast

from flask import g, request
from sqlalchemy import select

from dms.common.enums import EnabledStatus, RoleCode, SecurityLevel
from dms.common.errors import (
    AuthenticationError,
    AuthVersionMismatchError,
    ForbiddenError,
    ReauthRequiredError,
    TokenExpiredError,
    UserDisabledError,
)
from dms.extensions import db
from dms.models import User
from dms.security.auth_context import AuthContext, get_auth_context
from dms.security.jwt_tokens import (
    REAUTH_TOKEN_TYPE,
    decode_access_token,
    decode_reauth_token,
)

F = TypeVar("F", bound=Callable[..., Any])


def extract_bearer_token() -> str:
    authorization = request.headers.get("Authorization", "")
    scheme, separator, token = authorization.partition(" ")
    if not separator or scheme.lower() != "bearer" or not token.strip():
        raise AuthenticationError()
    return token.strip()


def authenticate_request(token: str) -> AuthContext:
    claims = decode_access_token(token)
    try:
        user_id = int(claims["sub"])
        token_auth_version = int(claims["authVersion"])
        token_username = str(claims["username"])
        token_role = str(claims["roleCode"])
    except (KeyError, TypeError, ValueError) as exc:
        raise AuthenticationError() from exc

    user = db.session.scalar(select(User).where(User.id == user_id))
    if user is None or user.is_deleted:
        raise AuthenticationError()
    if user.status != EnabledStatus.ENABLED:
        raise UserDisabledError()
    if user.auth_version != token_auth_version:
        raise AuthVersionMismatchError()
    if user.username != token_username or user.role_code != token_role:
        raise AuthenticationError()

    context = AuthContext(
        user_id=user.id,
        username=user.username,
        real_name=user.real_name,
        organization_id=user.organization_id,
        organization_name=user.organization_name,
        role_code=RoleCode(user.role_code),
        security_level=SecurityLevel(user.security_level),
        auth_version=user.auth_version,
    )
    g.dms_auth_context = context
    return context


def bearer_auth_required(view: F) -> F:
    @wraps(view)
    def wrapped(*args: Any, **kwargs: Any):
        authenticate_request(extract_bearer_token())
        return view(*args, **kwargs)

    return cast(F, wrapped)


def admin_required(view: F) -> F:
    @wraps(view)
    def wrapped(*args: Any, **kwargs: Any):
        if get_auth_context().role_code != RoleCode.ADMIN:
            raise ForbiddenError()
        return view(*args, **kwargs)

    return cast(F, wrapped)


def _require_manage_flag(can_manage) -> None:
    """加载当前用户并按管理位判定（admin 账号恒通过）。"""
    context = get_auth_context()
    user = db.session.get(User, context.user_id)
    if user is None or not can_manage(db.session, user):
        raise ForbiddenError()


def system_management_required(view: F) -> F:
    """系统管理页签接口鉴权（契约1.23）：manage_system 管理位。"""

    @wraps(view)
    def wrapped(*args: Any, **kwargs: Any):
        from dms.services.admin_permission import can_manage_system

        _require_manage_flag(can_manage_system)
        return view(*args, **kwargs)

    return cast(F, wrapped)


def audit_log_required(view: F) -> F:
    """日志审计页签接口鉴权（契约1.23）：manage_audit_log 管理位。"""

    @wraps(view)
    def wrapped(*args: Any, **kwargs: Any):
        from dms.services.admin_permission import can_view_audit_log

        _require_manage_flag(can_view_audit_log)
        return view(*args, **kwargs)

    return cast(F, wrapped)


def role_edit_required(view: F) -> F:
    """角色管理写操作鉴权（契约1.25）：manage_role_edit 管理位。"""

    @wraps(view)
    def wrapped(*args: Any, **kwargs: Any):
        from dms.services.admin_permission import can_edit_roles

        _require_manage_flag(can_edit_roles)
        return view(*args, **kwargs)

    return cast(F, wrapped)


REAUTH_HEADER = "X-DMS-Reauth-Token"


def require_reauth() -> None:
    """敏感操作二次认证：校验 X-DMS-Reauth-Token。

    缺失、无效、过期、非本人或 authVersion 不符一律拒绝（403 REAUTH_REQUIRED）。
    """
    context = get_auth_context()
    token = request.headers.get(REAUTH_HEADER, "").strip()
    if not token:
        raise ReauthRequiredError()
    try:
        claims = decode_reauth_token(token)
        if (
            str(claims.get("type")) != REAUTH_TOKEN_TYPE
            or int(claims["sub"]) != context.user_id
            or int(claims["authVersion"]) != context.auth_version
        ):
            raise ReauthRequiredError()
    except (
        AuthenticationError,
        TokenExpiredError,
        KeyError,
        TypeError,
        ValueError,
    ) as exc:
        raise ReauthRequiredError() from exc


def reauth_required(view: F) -> F:
    @wraps(view)
    def wrapped(*args: Any, **kwargs: Any):
        require_reauth()
        return view(*args, **kwargs)

    return cast(F, wrapped)