使用装饰器检查用户操作权限

4次阅读

共计 2266 个字符,预计需要花费 6 分钟才能阅读完成。

假设我们使用 flask 进行开发,需要检查当前用户是否具备某个操作权限。

from flask import abort
from flask_login import current_user
from .models import Permission

先从最简单的情况来想。f是某个操作函数,如添加评论。那么用户想要添加评论需要具备评论权限 Permission.COMMENT
检查操作权限最直白的逻辑当然是使用条件判断语句:

if not current_user.can(Permission.COMMENT):
    abort(403)
f(*args, **kwargs)

为了不在每一次调用 f 前都写一边 if 语句,我们将这一层逻辑包含到一个新的函数里,则有:

def decorated_function(*args, **kwargs):
    if not current_user.can(Permission.COMMENT):
        abort(403)
    return f(*args, **kwargs)

除非当前用户没有操作权限,decorated_function做的事情将与 f 做的事情一模一样。我们想要的就是一个这样的函数。
然而,目前 decorated_function 只能完成和 f 同样的事。如果对另一个函数 g 想要一个同样的函数,是否就要再写一个

def decorated_function_new(*args, **kwargs):
    if not current_user.can(Permission.COMMENT):
        abort(403)
    return g(*args, **kwargs)

了呢?

decorated_function_newdecorated_function 唯一的区别在于返回的值是 f(*args, **kwargs) 还是g(*args, **kwargs)。那么我们可以再定义一个函数,将 f 或者 g 作为参数传进去,就可以解决问题。

def decorator(f):
    def decorated_function(*args, **kwargs):
        if not current_user.can(Permission.COMMENT):
            abort(403)
        return f(*args, **kwargs)
    return decorated_function

这样,对于任何一个函数 f, 我们都能通过调用decorator(f) 的方式,得到一个与之对应的包含了在进行操作前检查评论权限的函数。注意这只是个函数,需要调用 (和传参) 才能得到返回值。

decorator(f)(*args, **kwargs)

这里,函数 decorator 就是一个装饰器,可以在不更改 f 定义的前提下在代码运行期间动态地增加 f 的功能。

@ 将这个装饰器放在某个函数定义之前,就可以起到替换原函数的作用。

@decorator
def add_comment():
    '''add a comment'''
    ...

这等价于:

add_comment = decorator(add_comment)

然而因为装饰器是替换原函数,我们失去了原函数的其他一些特征。比如:

print(add_comment.__name__)
print(add_comment.__doc__)

将输出 decorated_functionNone,而不是 add_commentadd a comment。如果用 decorator 装饰了多个函数,则它们的__name__和__doc__属性都会被掩盖掉,有些依赖函数签名的代码执行就会出错。为了避免这样的事情,可以使用 functools 库中的 wrapsf的相关属性复制进来。

from functools import wraps

...

def decorator(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if not current_user.can(Permission.COMMENT):
            abort(403)
        return f(*args, **kwargs)
    return decorated_function

现在我们已经实现了一个装饰器,但只能用它检查评论权限。我们可以定义一个装饰器工厂函数,产生检查各种权限的装饰器。

def permission_required(permission):
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if not current_user.can(permission):
                abort(403)
            return f(*args, **kwargs)
        return decorated_function
    return decorator

...

@permission_required(Permission.WRITE)
def new_post(user):
    ...

在这颗裹了三层 def 的洋葱中,最内层 def 添加了原函数之外的逻辑,第二层 deff作为参数,使得装饰器可以在其他任何函数上,第三层用于传入参数控制装饰器的行为。

如果我们希望用 permission_required 为检查管理员权限单独生成一个装饰器:

def admin_required(f):
    return permission_required(Permission.ADMIN)(f)

写成 admin_required = permission_required(Permission.ADMIN) 应该也可以,不过还是用定义函数的形式比较好吧。

正文完
 0