Files
windmill/gen.py
2023-11-24 15:17:59 +08:00

97 lines
3.0 KiB
Python

import os
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import hashes
import base64
import time
public_key_file = "public_key.pem"
private_key_file = "private_key.pem"
# 生成 RSA 公私钥对并缓存到文件
def generate_and_save_keys():
private_key = rsa.generate_private_key(public_exponent=65537, key_size=1024)
public_key = private_key.public_key()
# 将公钥序列化为 PEM 格式并保存
with open(public_key_file, "wb") as f:
f.write(
public_key.public_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PublicFormat.SubjectPublicKeyInfo,
)
)
# 将私钥序列化为 PEM 格式并保存
with open(private_key_file, "wb") as f:
f.write(
private_key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
)
)
return private_key
# 检查公私钥文件是否存在,如果不存在则生成新的
def get_or_generate_keys():
if not os.path.exists(public_key_file) or not os.path.exists(private_key_file):
return generate_and_save_keys()
with open(private_key_file, "rb") as f:
private_key = serialization.load_pem_private_key(
f.read(),
password=None,
)
return private_key
# 生成许可证密钥
def generate_license_key(client_id, expiry, private_key):
payload = f"{client_id}{expiry}"
signature = private_key.sign(
payload.encode(),
padding.PSS(
mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH
),
hashes.SHA256(),
)
signature_b64 = base64.b64encode(signature).decode()
license_key = f"{client_id}.{expiry}.{signature_b64}"
return license_key
# 示例用法
private_key = get_or_generate_keys()
client_id = "123"
expiry = int(time.time()) + 60 * 60 * 24 * 365 * 100 # 有效期为100年
license_key = generate_license_key(client_id, expiry, private_key)
public_key = private_key.public_key()
# 序列化公钥为 PKCS#8 DER 格式
der_public_key = public_key.public_bytes(
encoding=serialization.Encoding.DER,
format=serialization.PublicFormat.SubjectPublicKeyInfo,
)
# 将 DER 格式公钥编码为 Base64
b64_public_key = base64.b64encode(der_public_key).decode()
PREV_KEY = "MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQDgVShzcLSPiOi+8ET8fggob1kmi47/cE12JaidPkwfGnScZItghkqtiLsct0U4kJhlp5gO89DYTBmIKadvxwY7kMsLlZzmi2emVH7c27cByGASY8QmWDNdG4Ggy/NDflGGBdAtN6gHawZAg4zHv3qpbPQGHH1/6sXIohcXhOnouwIDAQAB"
print(len(b64_public_key), len(PREV_KEY))
assert len(b64_public_key) == len(PREV_KEY)
print("公钥:", b64_public_key)
print("许可证密钥:", license_key)
with open("license_key.txt", "w") as f:
f.write(license_key)