flask使用自定义装饰器鉴权

有的时候我们需要在自己的api里添加token校验来确保请求的安全性,那么我们就可以自己写一个装饰器来检查请求头里的授权令牌。
装饰器代码如下:

from functools import wraps
from flask import jsonify, request

def token_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        # access_token和BLACKLIST需要你自己定义引入,否则会报错
        BLACKLIST = set()
        access_token_headers = request.headers.get('Authorization', None)
        if not access_token_headers:
            return jsonify({"code": "err", "msg": "Missing authorization token"}), 401
        if access_token == None:
            return jsonify({"code": "err", "msg": "Please log in first!"}), 401
        if access_token_headers == access_token and access_token not in BLACKLIST:
            try:
                return f(*args, **kwargs)
            except Exception as e:
                return jsonify({"code": "异常", "message": "{}".format(e)})
        else:
            return jsonify({"msg": "The token is incorrect, please log in again!"}), 401
    return decorated_function

装饰器使用代码如下:
@app.route('/custom', methods=['POST'])
@token_required
def custom():
    printf("鉴权成功")

完整代码如下:
# 导入所需的库
from flask import Flask, abort, jsonify, request
from functools import wraps
from flask_jwt_extended import JWTManager, create_access_token
from werkzeug.security import generate_password_hash, check_password_hash

# 创建 Flask 应用实例
app = Flask(__name__)
access_token = None
global BLACKLIST
BLACKLIST = set()
# 生成密码哈希
password_hash = generate_password_hash('123456')
# 配置 JWT 密钥
app.config['JWT_SECRET_KEY'] = 'jwt-secret-string'
jwt = JWTManager(app)

# 定义 /login 路由,用于用户登录
@app.route('/login', methods=['POST'])
def login():
    # 获取请求参数
    password = request.json.get('password', None)
    if not check_password_hash(password_hash, password):
        # 密码错误
        abort(403)
    # 生成 access_token,并全局使用
    global access_token
    access_token = create_access_token(identity=password)
    return jsonify(code=200,access_token=access_token)

# 定义 /custom 路由,需要 token 验证
@app.route('/custom', methods=['POST'])
@token_required
def custom():
    print("鉴权成功")

# 定义 token_required 装饰器,用于验证 token
def token_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        # access_token 和 BLACKLIST 需要你自己定义引入,否则会报错
        access_token_headers = request.headers.get('Authorization', None)
        if not access_token_headers:
            return jsonify({"code": "err", "msg": "Missing authorization token"}), 401
        if access_token == None:
            return jsonify({"code": "err", "msg": "Please log in first!"}), 401
        if access_token_headers == access_token and access_token not in BLACKLIST:
            try:
                return f(*args, **kwargs)
            except Exception as e:
                return jsonify({"code": "异常", "message": "{}".format(e)})
        else:
            return jsonify({"msg": "The token is incorrect, please log in again!"}), 401
    return decorated_function

# 启动 Flask 应用
if __name__ == "__main__":
    app.run(debug=True, host='0.0.0.0', port=3000)

如果请求头中没有授权令牌,或者令牌不正确,那么装饰器将返回一个错误消息和401状态码,表示未经授权。
如果令牌正确,那么装饰器将执行被装饰的函数,并返回其结果。如果在执行过程中发生异常,装饰器将捕获这个异常,并返回一个包含异常信息的错误消息。