"""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) def auditor_required(view: F) -> F: @wraps(view) def wrapped(*args: Any, **kwargs: Any): if get_auth_context().role_code != RoleCode.AUDITOR: raise ForbiddenError() return view(*args, **kwargs) return cast(F, wrapped)