finance_data_process/test_token.py

"""生成测试token
"""

import base64  # 导入base64模块，用于编码和解码数据
import json  # 导入json模块，用于处理JSON格式数据
from Crypto.Cipher import AES  # 导入AES加密模块
from datetime import datetime  # 导入datetime模块，用于获取当前时间

# 定义加密密钥
key = "2na$$PdV9AW8b#CS"
# 定义一个常量字符串，用于替换加密后的Base64字符串中的特殊字符“+”
ADD = "/add/"


def encrypt(str, key):
    """
    使用AES算法对输入字符串进行加密，并返回加密后的Base64编码字符串。

    参数:
        str (str): 要加密的字符串
        key (str): 加密密钥

    返回:
        str: 加密后的Base64编码字符串
    """
    # 创建AES加密器，使用ECB模式（电子密码本模式）
    cipher = AES.new(key.encode('utf-8'), AES.MODE_ECB)
    # 对输入字符串进行填充，使其长度符合AES块大小的要求
    padded_data = pad(str.encode('utf-8'), AES.block_size)
    # 使用AES加密器对填充后的数据进行加密
    encrypted = cipher.encrypt(padded_data)
    # 将加密后的二进制数据转换为Base64编码字符串
    encrypted_base64 = base64.b64encode(encrypted).decode('utf-8')
    # 去除Base64编码字符串中的换行符（如果有）
    encrypted_base64 = encrypted_base64.replace('\n', '')
    # 替换Base64编码字符串中的“+”符号为自定义字符串“add”，避免特殊字符冲突
    encrypted_base64 = encrypted_base64.replace('+', ADD)
    return encrypted_base64


def pad(data, block_size):
    """
    对输入数据进行填充，使其长度成为块大小的整数倍（PKCS7填充方式）。

    参数:
        data (bytes): 要填充的字节数据
        block_size (int): 块大小（AES块大小为16字节）

    返回:
        bytes: 填充后的字节数据
    """
    # 计算需要填充的字节数
    padding = block_size - len(data) % block_size
    # 按照PKCS7填充方式，在数据末尾添加padding个值为padding的字节
    return data + bytes([padding]) * padding


# 构造一个包含必要信息的字典
s = {
    "APP_ID": "FINANCE",  # 应用ID
    "TIME_STAMP": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),  # 当前时间戳
    "REQUEST_URL": "",  # 请求URL（此处为空字符串）
    "LOGIN_ID": "test",  # 登录ID
    "EXPIRE_TIME": 3600,  # 过期时间（单位：秒）
}

# 将字典转换为JSON格式字符串，并调用encrypt函数对其进行加密最后打印加密结果
print(encrypt(json.dumps(s), key))