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,
    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 decode_access_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)