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"{payload}.{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)

REPLACE_TEMPLATE = f'''
public_key = b"""{PREV_KEY}"""

new_public_key = b"""{b64_public_key}"""

from sys import argv

with open(argv[1], "rb") as f:
    data = f.read()

assert public_key in data

with open(argv[1], "wb") as f:
    f.write(data.replace(public_key, new_public_key))
'''

with open("replace.py", "w") as f:
    f.write(REPLACE_TEMPLATE)

print("许可证密钥:", license_key)

with open("license_key.txt", "w") as f:
    f.write(license_key)