【flask】jwt登录认证（有刷新token）

港控/mmm° 2022-12-19 14:27 717阅读 0赞

#### 1.获取token ####

每次请求都会获取token，token位于客户端的header的Authorization字段，取不到返回None  
common/utils/middleware.py

# common/utils/middleware.py
    from flask import request, current_app, g
    from utils.jwt_util import verify_jwt
    
    def get_userinfo():
        # 1.获取请求头中的token信息 [登录token，刷新token]
        # 规定：前端通过Authorization字段携带token
        token = request.headers.get("Authorization")
        # 获取秘钥
        key = current_app.config["JWT_SECRET"]
        # 2.进行token校验
        # 载荷信息有可能为空：token为空，token错误，token过期
        payload = verify_jwt(token, key)
        # 默认值
        g.user_id = None
        g.is_refresh = False
        if payload:
            # 3.从载荷字典中提取用户身份信息
            g.user_id = payload.get("user_id")
            g.is_refresh = payload.get("is_refresh", False)

使用钩子函数，每次请求都会尝试提取token  
app/\_\_init\_\_.py

# app/__init__.py
    def register_extensions(app):
        """组件初始化"""
        ...
        # 添加请求钩子
        from utils.middleware import get_userinfo
        app.before_request(get_userinfo)

#### 2.登录装饰器 ####

定义登录装饰器，哪里要求登录就调用这个装饰器  
common/utils/decorators.py

# common/utils/decorators.py
    from flask import g
    
    # 强制登录装饰器
    def login_reqiured(view_func):
        def wrapper(*args, **kwargs):
    
            user_id = g.user_id
            is_refresh = g.is_refresh
            # 1.判断用户是否登录
            if user_id is None:
                # 用户未登录
                # 401 权限认证失败
                return { "message": "invalid token"}, 401
            # 2.判断是否是刷新token
            elif user_id is not None and is_refresh:
                # 刷新token
                return { "message": "refresh token not for login"}, 403
            else:
                # 登录token
                return view_func(*args, **kwargs)
    
        return wrapper

#### 3.主要业务 ####

app/resources/user/passport.py

# app/resources/user/passport.py
    import random
    
    from flask_restful import Resource
    from flask_restful.reqparse import RequestParser
    
    from app import redis_client
    from utils.constants import SMS_CODE_EXPIRE
    
    from datetime import datetime, timedelta
    from flask import current_app, g
    from flask_restful.inputs import regex
    from flask_restful.reqparse import RequestParser
    from sqlalchemy.orm import load_only
    from app import db
    from utils.parser import mobile as mobile_type
    from models.user import User
    from utils.jwt_util import generate_jwt
    from utils.decorators import login_required
    
    
    class SMSCodeResource(Resource):
        def get(self, mobile):
            rand_num = "%06d" % random.randint(0, 999999)
            key = "app:code:{}".format(mobile)
            redis_client.set(key, rand_num, ex=SMS_CODE_EXPIRE)
            print(f"短信验证码：mobile:{mobile},code:{rand_num}")
            return { "mobile": mobile}
    
    
    class LoginResource(Resource):
        # 内部生成token的方法
        def _generator_token(self, user_id):
            login_payload = { 
                "user_id": user_id,
                "is_refresh": False
            }
            expire_2h = datetime.utcnow() + timedelta(hours=current_app.config["LOGIN_TOKEN_EXPIRE"])
            secret = current_app.config["JWT_SECRET"]
            login_token = generate_jwt(payload=login_payload, expiry=expire_2h, secret=secret)
    
            refresh_payload = { 
                "user_id": user_id,
                "is_refresh": True
            }
            expire_14d = datetime.utcnow() + timedelta(days=current_app.config["REFRESH_TOKEN_EXPIRE"])
            refresh_token = generate_jwt(payload=refresh_payload, expiry=expire_14d, secret=secret)
            return login_token, refresh_token
    
        # 注册、登录
        def post(self):
            # 获取参数
            parser = RequestParser()
            parser.add_argument('mobile', required=True, location='json', type=mobile_type)
            parser.add_argument('code', required=True, location='json', type=regex(r'^\d{6}$'))
            args = parser.parse_args()
            mobile = args.mobile
            code = args.code
    
            # 校验短信验证码
            key = 'app:code:{}'.format(mobile)
            real_code = redis_client.get(key)
            if not real_code or real_code != code:
                return { 'message': 'Invalid Code', 'data': None}, 400
            # 删除验证码
            # redis_client.delete(key)
    
            # 校验成功, 查询数据库
            user = User.query.options(load_only(User.id)).filter(User.mobile == mobile).first()
            if user:  # 如果有, 取出用户id, 更新最后登录时间
                user.last_login = datetime.now()
            else:  # 如果没有, 创建新用户
                user = User(mobile=mobile, name=mobile, last_login=datetime.now())
                db.session.add(user)
            try:
                db.session.commit()
            except Exception as e:
                db.session.rollback()
                return { "message": e}, 507
            # 生成jwt
            login_token, refresh_token = self._generator_token(user.id)
            # 返回结果
            return { "login_token": login_token, "refresh_token": refresh_token}
    
        # 刷新token
        def put(self):
            user_id = g.user_id
            is_refresh = g.is_refresh
            if user_id and is_refresh:
                login_token, _ = self._generator_token(user_id)
                return { "new_token": login_token}
            else:
                return { "message": "wrong refresh token"}, 403
    
    
    # 获取当前登录用户信息类视图
    class CurrentUserResource(Resource):
        method_decorators = { 
            "get": [login_required]
        }
    
        def get(self):
            user_id = g.user_id
            user = User.query.options(load_only(
                User.id,
                User.name,
                User.profile_photo,
                User.introduction,
                User.article_count,
                User.following_count,
                User.fans_count
            )).filter(User.id == user_id).first()
            user_dict = user.to_dict()
            return user_dict

common/utils/jwt\_util.py

# common/utils/jwt_util.py
    import jwt
    from flask import current_app
    
    
    def generate_jwt(payload, expiry, secret=None):
        """ 生成jwt :param payload: dict 载荷 :param expiry: datetime 有效期 :param secret: 密钥 :return: jwt """
        _payload = { 'exp': expiry}
        _payload.update(payload)
    
        if not secret:
            secret = current_app.config['JWT_SECRET']
    
        token = jwt.encode(_payload, secret, algorithm='HS256')
        return token.decode()
    
    
    def verify_jwt(token, secret=None):
        """ 检验jwt :param token: jwt :param secret: 密钥 :return: dict: payload """
        if not secret:
            secret = current_app.config['JWT_SECRET']
    
        try:
            payload = jwt.decode(token, secret, algorithm=['HS256'])
        except jwt.PyJWTError:
            payload = None
    
        return payload

common/utils/parser.py

def mobile(mobile_str):
        """ 检验手机号格式 :param mobile_str: str 被检验字符串 :return: mobile_str """
        if re.match(r'^1[3-9]\d{9}$', mobile_str):
            return mobile_str
        else:
            raise ValueError('{} is not a valid mobile'.format(mobile_str))