59 lines
1.6 KiB
Python
59 lines
1.6 KiB
Python
# -*- coding: utf-8 -*-
|
|
# @Time : 2022/5/14 16:19
|
|
# @Author : 臧成龙
|
|
# @FileName: fu_jwt.py
|
|
# @Software: PyCharm
|
|
from typing import Union
|
|
from simplejwt import util
|
|
from simplejwt.jwt import default_alg, _hash, Jwt
|
|
import datetime
|
|
import json
|
|
|
|
|
|
class DateEncoder(json.JSONEncoder):
|
|
def default(self, obj):
|
|
if isinstance(obj, datetime.datetime):
|
|
return obj.strftime("%Y-%m-%d %H:%M:%S")
|
|
else:
|
|
return json.JSONEncoder.default(self, obj)
|
|
|
|
|
|
class FuJwt(Jwt):
|
|
|
|
def encode(self) -> str:
|
|
payload = {}
|
|
payload.update(self.registered_claims)
|
|
payload.update(self.payload)
|
|
return encode(self.secret, payload, self.alg, self.header)
|
|
|
|
|
|
def encode(secret: Union[str, bytes], payload: dict = None, alg: str = default_alg, header: dict = None) -> str:
|
|
"""
|
|
:param secret: The secret used to encode the token.
|
|
:type secret: Union[str, bytes]
|
|
:param payload: The payload to be encoded in the token.
|
|
:type payload: dict
|
|
:param alg: The algorithm used to hash the token.
|
|
:type alg: str
|
|
:param header: The header to be encoded in the token.
|
|
:type header: dict
|
|
:return: A new token
|
|
:rtype: str
|
|
"""
|
|
secret = util.to_bytes(secret)
|
|
|
|
payload = payload or {}
|
|
header = header or {}
|
|
|
|
header_json = util.to_bytes(json.dumps(header))
|
|
header_b64 = util.b64_encode(header_json)
|
|
payload_json = util.to_bytes(json.dumps(payload, cls=DateEncoder))
|
|
payload_b64 = util.b64_encode(payload_json)
|
|
|
|
pre_signature = util.join(header_b64, payload_b64)
|
|
signature = _hash(secret, pre_signature, alg)
|
|
signature_b64 = util.b64_encode(signature)
|
|
|
|
token = util.join(pre_signature, signature_b64)
|
|
return util.from_bytes(token)
|