Python实现Token过期自动刷新并重试原请求

问题描述

本人系统经常与上游系统进行交互，其中获取 Token 的接口调用次数过多，想将它存储在内存中，在请求失败后再重新获取 Token 继续进行原请求。

解决方案

灵感来自 sbzhu 的项目 weworkapi_python，十分巧妙

def httpCall(self, urlType, args=None) : 
    shortUrl = urlType[0]
    method = urlType[1]
    response = {}
    for retryCnt in range(0, 3) :  # 重试三次，可用于重新获取 Token 后继续进行原请求
    # 正常进行原请求
        if 'POST' == method :
            url = self.__makeUrl(shortUrl)
            response = self.__httpPost(url, args)
        elif 'GET' == method :
            url = self.__makeUrl(shortUrl)
            url = self.__appendArgs(url, args)
            response = self.__httpGet(url)
        else : 
            raise ApiException(-1, "unknown method type")

        # check if token expired
        if self.__tokenExpired(response.get('errcode')) :  # 校验 Token 是否过期
            self.__refreshToken(shortUrl)  # 过期则刷新
            retryCnt += 1
            continue
        else :  # 不过期则跳出重试机制
            break

    return self.__checkResponse(response) 

代码

上游服务器

import os
import datetime
import traceback

import jwt
from flask import Flask, request, current_app, jsonify, abort

app = Flask(__name__)
app.config['SECRET_KEY'] = os.environ.get('SECRET_KEY') or 'this is a secret'
pass_paths = ['/get_token']  # 跳过验证的访问路径


@app.before_request
def before_request():
    """Token验证"""
    need_verify = True  # 是否需要验证
    for path in pass_paths:
        if request.path.endswith(path):
            need_verify = False
            break

    if need_verify:
        try:
            Token = request.headers.get('Token')
            key = current_app.config['SECRET_KEY']
            jwt.decode(Token, key, algorithms=['HS256'])
        except:
            traceback.print_exc()
            abort(401)


@app.route('/get_token', methods=['POST'])
def get_token():
    """获取Token"""
    username = request.form.get('username')
    password = request.form.get('password')
    if username == 'admin' and password == '123456':  # 身份校验
        # exp = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(minutes=1)  # 一分钟后
        exp = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(seconds=3)  # 3秒后
        payload = {'username': username, 'exp': exp}
        key = current_app.config['SECRET_KEY']
        Token = jwt.encode(payload, key, algorithm='HS256')
        return jsonify({'Token': Token})
    else:
        abort(401)


@app.route('/get_datetime')
def get_datetime():
    return jsonify({'time': datetime.datetime.now()})


if __name__ == '__main__':
    app.run()

用户名密码简略为 admin 和 123456

下游服务器

import time

import requests

session = requests.Session()


def raise_exception(response, *args, **kwargs):
    if response.status_code == 401:
        print('Retry')
        headers = get_headers()
        response.request.headers.update(headers)
        return session.send(response.request, verify=False)


hooks = {'response': raise_exception}


def get_headers():
    get_token_api = 'http://127.0.0.1:5000/get_token'
    data = {'username': 'admin', 'password': '123456'}
    response = requests.post(get_token_api, data=data)
    Token = response.json()['Token']
    headers = {'Token': Token}
    return headers


get_datetime_api = 'http://127.0.0.1:5000/get_datetime'
headers = get_headers()
time.sleep(4)  # 模拟Token过期
response = requests.get(get_datetime_api, headers=headers, hooks=hooks)
print(response.json())

参考文献

1. weworkapi_python GitHub
2. Python requests 高级用法：timeouts、retries、hooks
3. Python Requests - retry request after re-authentication
4. RESTful Authentication with Flask