d3i-szct/paste-framework/paste/web/decorators.py

import functools
import logging
from typing import Awaitable

from jwt import ExpiredSignatureError, InvalidSignatureError, InvalidTokenError

from paste.security import token
from paste.web.handler import RequestHandler


def route(route_pattern: str):
    """
    路由装饰器。为类增加 route_pattern 属性，并赋值。

    :param route_pattern: URL 路径模式
    """

    def wrapper(cls: type[RequestHandler]):
        cls.route_pattern = route_pattern
        return cls

    # 标记已经被 route 装饰
    setattr(wrapper, '__route__', True)
    return wrapper


def auth_token(func):
    """
    令牌验证装饰器，用于 :class:`tornado.web.RequestHandler` 子类中的 get()/post() 等需要执行权限验证的方法，以便在正
    式执行方法前利用客户端提交的令牌进行鉴权。

    当执行该装饰器解码令牌后，将更新 RequestHandler 对象的 token_payload 属性数据。其次，若能通过令牌中配载的 user_id 取
    得用户数据，则还将设置 current_user 属性。

    该装饰器仅用于校验令牌的有效性，并取得用户信息，不负责校验用户的具体权限，若要验证权限需使用 @auth_permission 装饰器。

    使用方式如下::

        @auth_token
        async def post(self):
            pass

    要求在请求的 Headers 中必须包含 Access-Token，且内容由 encode_token 方法签发。

    :param func: 被装饰的函数对象，不需要手动传入该参数
    :return: 装饰后的函数对象
    """

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        req_handler: RequestHandler = args[0]
        try:
            # 请求头
            req_headers = dict(req_handler.request.headers)

            # 取出 Token
            access_token = req_headers.get('Access-Token', None)
            if access_token in (None, b'', ''):
                raise InvalidTokenError(f'请求地址：{req_handler.request.uri}')

            # 如果采用 OAuth2 规范，这里应当调用远程 API 执行解码，解码后返回用户信息

            # 用解码后的 Token 字典更新 Handler 中的的 token_dict
            token_payload = token.decode_token(access_token)
            req_handler.token_payload.update(token_payload)

            # 根据 Token 读取用户对象，并设置到请求处理对象（控制器）
            _user_id = req_handler.token_param('user_id')
            if _user_id and req_handler.user_class:
                _user = await req_handler.user_class.async_find_by_id(_user_id)
                if _user is None:
                    raise InvalidTokenError()
                req_handler.current_user = _user
                await req_handler.after_auth_token(token_payload)

            # 兼容同步或异步方法
            _result = func(*args, **kwargs)
            if isinstance(_result, Awaitable):
                _result = await _result
            return _result
        except ExpiredSignatureError as e:
            e.args = ('令牌已过期，请求被拒绝.',)
            req_handler.response_error(e, status_code=403, api_status_code=403)
            req_handler.log(msg=e, level=logging.ERROR, is_log_exc=True)
            return None
        except InvalidSignatureError as e:
            e.args = ('令牌签名错误，请求被拒绝.',)
            req_handler.response_error(e, status_code=403, api_status_code=403)
            req_handler.log(msg=e, level=logging.ERROR, is_log_exc=True)
            return None
        except InvalidTokenError as e:
            e.args = ('令牌错误，请求被拒绝.',)
            req_handler.response_error(e, status_code=401, api_status_code=401)
            req_handler.log(msg=e, level=logging.ERROR, is_log_exc=True)
            return None
        except Exception as e:
            req_handler.response_error(e, status_code=501, api_status_code=501)
            req_handler.log(msg=e, level=logging.ERROR, is_log_exc=True)
            return None

    # 标记已经被 auth_token 装饰
    setattr(wrapper, '__auth_token__', True)
    return wrapper


def auth_permission(func):
    """
    权限检查装饰器。若不启用 RBAC 则不应用该装饰器。
    用于检查用户是否有执行某个操作的具体权限。该装饰器须跟随在 @auth_token 装饰器的后面使用。

    :param func: 被装饰的函数对象，不需要手动传入该参数
    :return: 装饰后的函数对象
    """

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        #
        # 为了能在不使用 RBAC 的系统中正常运行，这里的引用必须放在函数中
        # 否则初始化过程中 RBAC 数据模型会尝试读取数据表配置，发生错误
        #
        from paste.rbac.rbac_user import RbacUser, Supervisors

        req_handler: RequestHandler = args[0]
        try:
            # 验证当前用户状态
            _user: RbacUser = req_handler.current_user
            assert _user is not None, f"无效令牌或未登录，无权执行：{req_handler.route_pattern} 操作."

            # 类型检测
            _right_type = isinstance(_user, RbacUser)
            assert _right_type, f"当前用户类型错误，必须为 RbacUser 的子类."

            if _user.username not in Supervisors:
                # 验证用户权限状态
                _has_permission = await _user.has_permission(req_handler.route_pattern)
                assert _has_permission, f"当前用户 {_user.username} 无权执行：{req_handler.route_pattern} 操作."

            # 兼容同步或异步方法
            _result = func(*args, **kwargs)
            if isinstance(_result, Awaitable):
                _result = await _result
            return _result
        except AssertionError as e:
            req_handler.response_error(e, status_code=401, api_status_code=401)
            req_handler.log(msg=e, level=logging.ERROR, is_log_exc=True)
            return None
        except Exception as e:
            req_handler.response_error(e, status_code=501, api_status_code=501)
            req_handler.log(msg=e, level=logging.ERROR, is_log_exc=True)
            return None

    # 标记已经被 auth_permission 装饰
    setattr(wrapper, '__auth_permission__', True)
    return wrapper


def auth_rule(func):
    """
    规则检查装饰器。若不启用规则验证，则不应用该装饰器。
    用于对用户按规则验证。该装饰器须跟随在 @auth_token 装饰器的后面使用。

    :param func: 被装饰的函数对象，不需要手动传入该参数
    :return: 装饰后的函数对象
    """

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        #
        # 为了能在不使用 RBAC 的系统中正常运行，这里的引用必须放在函数中
        # 否则初始化过程中 RBAC 数据模型会尝试读取数据表配置，发生错误
        #
        from paste.rbac.rbac_user import RbacUser, Supervisors

        req_handler: RequestHandler = args[0]
        try:
            # 验证当前用户状态
            _user: RbacUser = req_handler.current_user
            assert _user is not None, f"无效令牌或未登录，无权执行：{req_handler.route_pattern} 操作."

            # 类型检测
            _right_type = isinstance(_user, RbacUser)
            assert _right_type, f"当前用户类型错误，必须为 RbacUser 的子类."

            if _user.username not in Supervisors:
                # 验证用户权限状态
                _user_can = await _user.can(req_handler.route_pattern, **kwargs)
                assert _user_can, f"当前用户  {_user.username} 无权执行：{req_handler.route_pattern} 操作（规则验证不通过）."

            # 兼容同步或异步方法
            _result = func(*args, **kwargs)
            if isinstance(_result, Awaitable):
                _result = await _result
            return _result
        except AssertionError as e:
            req_handler.response_error(e, status_code=401, api_status_code=401)
            req_handler.log(msg=e, level=logging.ERROR, is_log_exc=True)
            return None
        except Exception as e:
            req_handler.response_error(e, status_code=501, api_status_code=501)
            req_handler.log(msg=e, level=logging.ERROR, is_log_exc=True)
            return None

    # 标记已经被 auth_rule 装饰
    setattr(wrapper, '__auth_rule__', True)
    return wrapper