import copy
import json
import base64
import hmac
import time


class Jwt:
    def __init__(self):
        pass

    @staticmethod
    def encode(payload, key, exp=300):
        header = {
            'alg': 'HS256',
            'typ': 'JWT',
        }
        # separators控制字典转为json后键值对什么相连接
        # sort_keys让json窜出来为有序的
        header = json.dumps(header, separators=(',', ':'), sort_keys=True)
        header = Jwt.b64encode(header.encode())

        payload = copy.deepcopy(payload)
        payload['exp'] = int(time.time() + exp)
        payload = json.dumps(payload, separators=(',', ':'), sort_keys=True)
        payload = payload.encode()
        payload = Jwt.b64encode(payload)
        sign = header + b'.' + payload

        # 判断key的类型，如果不是byte，转换为byte
        if isinstance(key, str):
            key = key.encode()
        h = hmac.new(key, sign, digestmod='SHA256')
        sign = h.digest()
        sign = Jwt.b64encode(sign)
        return header + b'.' + payload + b'.' + sign

    @staticmethod
    def b64encode(s):
        # 替换原生base64中的等号为空
        return base64.urlsafe_b64encode(s).replace(b"=", b"")

    @staticmethod
    def b64decode(bs):
        # 将替换=后的base64补回至原长度
        length = len(bs) % 4
        length = 4 - length
        bs = bs + b"=" * length
        return base64.urlsafe_b64decode(bs)

    @staticmethod
    def decode(token, key):
        """
        校验token
        :param token:
        :param key:
        :return:
        """
        header_bs, payload_bs, sign = token.split(b'.')
        if isinstance(key, str):
            key = key.encode()
        h = hmac.new(key, header_bs + b'.' + payload_bs, digestmod='SHA256')
        new_token = Jwt.b64encode(h.digest())
        if sign == new_token:
            # token合法判断是否过期
            payload_j = Jwt.b64decode(payload_bs)
            payload = json.loads(payload_j)
            # 过期时间戳
            exp = payload['exp']
            now = time.time()
            # 对比是否过期
            if now > exp:
                # token过期
                raise JwtError('Your token is expired')
            return payload

        else:
            # 当前传过来的token违法，则raise
            raise JwtError('Your token is valid')
        # 如果校验成功
        # return payload 明文
        # 如果没有校验成功
        # raise具体错误
        # 校验成功标准：1：签名相符，2：是否过期


class JwtError(Exception):
    """
    自定义异常
    """

    def __init__(self, error_msg):
        self.error = error_msg

    def __str__(self):
        return '<JwtError error> %s' % self.error


if __name__ == "__main__":
    j = Jwt()
    print(j.decode(j.encode({"username": "guoxiaonaos"}, 'abcdef1234', 100), 'abcdef1234'))