135 lines
4.3 KiB
Python
135 lines
4.3 KiB
Python
import datetime
|
||
from functools import wraps
|
||
|
||
from flask import g, jsonify, request
|
||
|
||
import jwt
|
||
from app.configs.default import SECRET_KEY
|
||
from flask_httpauth import HTTPTokenAuth, HTTPAuth
|
||
|
||
# 生成token,有效时间为 60*60 秒
|
||
from app.exts import db
|
||
|
||
|
||
def generate_auth_token(user_name, expiration=3600):
|
||
reset_token = jwt.encode(
|
||
{
|
||
"user_name": user_name,
|
||
"exp": datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(seconds=expiration)
|
||
},
|
||
SECRET_KEY,
|
||
algorithm="HS256"
|
||
)
|
||
return reset_token
|
||
|
||
|
||
# 重写 HTTPTokenAuth
|
||
class HTTPTokenAuthReReturn(HTTPTokenAuth):
|
||
def __init__(self, scheme=None, realm=None, header=None):
|
||
super().__init__(scheme, realm, header)
|
||
|
||
# 重写default_auth_error, 或者也可以重新定义一个函数
|
||
# def default_auth_error(status):
|
||
# return jsonify(data={}, message="token 错误！", code=status), status
|
||
|
||
def default_auth_error(status, message):
|
||
return jsonify(data={}, message=message, code=status), status
|
||
|
||
# 如果重新定义函数的话，这里就传入新定义的函数名
|
||
super().error_handler(default_auth_error)
|
||
# 重写 login_required
|
||
|
||
def login_required(self, f=None, role=None, optional=None):
|
||
if f is not None and \
|
||
(role is not None or optional is not None): # pragma: no cover
|
||
raise ValueError(
|
||
'role and optional are the only supported arguments')
|
||
|
||
def login_required_internal(f):
|
||
@wraps(f)
|
||
def decorated(*args, **kwargs):
|
||
auth = self.get_auth()
|
||
|
||
if request.method != 'OPTIONS':
|
||
password = self.get_auth_password(auth)
|
||
|
||
status = None # 添加状态信息
|
||
message = None
|
||
user = self.authenticate(auth, password)
|
||
# 这里判断verify_token的返回值是否是这里的一员
|
||
if user in (False, None, 'BadSignature', 'SignatureExpired'):
|
||
status = 401
|
||
if user == 'BadSignature':
|
||
message = "Bad Signature"
|
||
elif user == 'SignatureExpired':
|
||
message = "Signature Expired"
|
||
elif not self.authorize(role, user, auth):
|
||
status = 403
|
||
message = "Forbidden"
|
||
if not optional and status:
|
||
# Clear TCP receive buffer of any pending data
|
||
request.data
|
||
try:
|
||
# 因为之前重写了default_auth_error所以多传入一个message
|
||
return self.auth_error_callback(status, message)
|
||
except TypeError:
|
||
return self.auth_error_callback()
|
||
|
||
g.flask_httpauth_user = user if user is not True \
|
||
else auth.username if auth else None
|
||
return f(*args, **kwargs)
|
||
|
||
return decorated
|
||
|
||
if f:
|
||
return login_required_internal(f)
|
||
return login_required_internal
|
||
|
||
|
||
auth = HTTPTokenAuthReReturn(scheme="Bearer")
|
||
|
||
|
||
# 获取用户权限
|
||
@auth.get_user_roles
|
||
def get_user_roles(user):
|
||
the_role = get_role_token(user["token"])
|
||
return the_role.role_key
|
||
|
||
|
||
# 验证token
|
||
@auth.verify_token
|
||
def verify_token(token):
|
||
try:
|
||
print(token)
|
||
data = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
|
||
print(data)
|
||
except jwt.ExpiredSignatureError:
|
||
print("token过期")
|
||
# 这里不用False，而是用自定义字符串
|
||
return "SignatureExpired"
|
||
except jwt.PyJWTError:
|
||
print("token错误")
|
||
# 这里不用False，而是用自定义字符串
|
||
return "BadSignature"
|
||
return True
|
||
|
||
|
||
# 解析token不加密部分
|
||
def token_parse(token):
|
||
the_token = str.replace(str(token), 'Bearer ', '')
|
||
decoded = jwt.decode(the_token, options={"verify_signature": False})
|
||
return decoded
|
||
|
||
|
||
# 根据 username 获取角色
|
||
def get_role_username(username: str):
|
||
return []
|
||
|
||
|
||
# 根据 username 获取角色
|
||
def get_role_token(token: str):
|
||
user_info = token_parse(token)
|
||
username = user_info["user_name"]
|
||
the_role = get_role_username(username)
|
||
return the_role
|