import base64
import json
from Crypto.Signature import PKCS1_v1_5
from Crypto.PublicKey import RSA
from Crypto.Hash import SHA512

REPLACE_TEMPLATE = '''
public_key = b"""-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAyZmShlU8Z8HdG0IWSZ8r
tSyzyxrXkJjsFUf0Ke7bm/TLtIggRdqOcUF3XEWqQk5RGD5vuq7Rlg1zZqMEBk8N
EZeRhkxyaZW8pLjxwuBUOnXfJew31+gsTNdKZzRjrvPumKr3EtkleuoxNdoatu4E
HrKmR/4Yi71EqAvkhk7ZjQFuF0osSWJMEEGGCSUYQnTEqUzcZSh1BhVpkIkeu8Kk
1wCtptODixvEujgqVe+SrE3UlZjBmPjC/CL+3cYmufpSNgcEJm2mwsdaXp2OPpfn
a0v85XL6i9ote2P+fLZ3wX9EoioHzgdgB7arOxY50QRJO7OyCqpKFKv6lRWTXuSt
hwIDAQAB
-----END PUBLIC KEY-----"""

new_public_key = b"""<NEW_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))
'''


def create_rsa_pair(is_save=False):
    '''
    创建rsa公钥私钥对
    :param is_save: default:False
    :return: public_key, private_key
    '''
    f = RSA.generate(2048)
    private_key = f.exportKey("PEM")  # 生成私钥
    public_key = f.publickey().exportKey()  # 生成公钥
    if is_save:
        with open("private.pem", "wb") as f:
            f.write(private_key)
        with open("public.pem", "wb") as f:
            f.write(public_key)
    return public_key, private_key


data = {
    "id": "j4cj88c98pnkj8j3qrzzngsijo",
    "issued_at": 1668153964204,
    "starts_at": 1668125164204,
    "expires_at": 2670745964204,
    "sku_name": "Self-Hosted Enterprise",
    "sku_short_name": "enterprise",
    "customer": {
        "id": "j4cj88c98pnkj8j3qrzzngsijo",
        "name": "114514",
        "email": "1145141919810",
        "company": "1145141919810.baidu.com"
    },
    "features": {
        "users": 114514,
        "ldap": True,
        "ldap_groups": True,
        "mfa": True,
        "google_oauth": True,
        "office365_oauth": True,
        "compliance": True,
        "cluster": True,
        "metrics": True,
        "mhpns": True,
        "saml": True,
        "elastic_search": True,
        "announcement": True,
        "theme_management": True,
        "email_notification_contents": True,
        "data_retention": True,
        "message_export": True,
        "custom_permissions_schemes": True,
        "custom_terms_of_service": True,
        "guest_accounts": True,
        "guest_accounts_permissions": True,
        "id_loaded": True,
        "lock_teammate_name_display": True,
        "cloud": False,
        "shared_channels": True,
        "remote_cluster_service": True,
        "openid": True,
        "enterprise_plugins": True,
        "advanced_logging": True,
        "future_features": True
    },
    "is_trial": False,
    "is_gov_sku": False
}

# public_key, private_key = create_rsa_pair(is_save=True)

with open("private.pem", "rb") as f:
    private_key = f.read()

with open("public.pem", "rb") as f:
    public_key = f.read()

plain_text = json.dumps(data)

print(plain_text)

signature = PKCS1_v1_5.new(RSA.importKey(private_key)).sign(SHA512.new(plain_text.encode('utf-8')))

with open("lic.txt", "wb") as f:
    f.write(base64.b64encode(plain_text.encode() + signature))

with open("replace.py", "w") as f:
    f.write(REPLACE_TEMPLATE.replace("<NEW_KEY>", public_key.decode()))