import os
import time
import json
import base64
from Crypto.Hash import SHA256
from Crypto.PublicKey import RSA
from Crypto.Signature import PKCS1_v1_5


class RSAK:
    @staticmethod
    def keytext(pempath: str) -> str:
        """
        读取 RSA PEM 文件，获取 RSA KEY 文本。
        """
        with open(pempath, mode="r", encoding="utf-8") as f:
            content = f.read().strip()
        return content

    @staticmethod
    def rsakey(keytext: str) -> RSA.RsaKey:
        """
            将 RSA密钥文本 转成 RSA密钥对象
        """
        return RSA.import_key(keytext.encode("utf-8"))

    @staticmethod
    def rsakeyfromfile(pempath: str) -> RSA.RsaKey:
        """
            读取 RSA PEM 文件, 获取 RSA密钥对象
        """
        keytext = RSAK.keytext(pempath)
        return RSAK.rsakey(keytext)

    @staticmethod
    def signature(data: dict, prikey: RSA.RsaKey):
        """
            计算签名值
        """
        prikeypkcs = PKCS1_v1_5.new(prikey)
        databytes = json.dumps(data, separators=(',', ':'), ensure_ascii=False).encode('utf-8')
        datahash = SHA256.new(databytes)
        return base64.b64encode(prikeypkcs.sign(datahash)).decode('utf-8')

    @staticmethod
    def verify(data: dict, signature: str, pubkey: RSA.RsaKey):
        """
            签名值验证
        """
        pubkeypkcs = PKCS1_v1_5.new(pubkey)
        signaturebytes = base64.b64decode(signature.encode('utf-8'))
        databytes = json.dumps(data, separators=(',', ':'), ensure_ascii=False).encode('utf-8')
        datahash = SHA256.new(databytes)
        return pubkeypkcs.verify(datahash, signaturebytes)
        
    @staticmethod
    def genRSAKey(dir: str, key_size: int = 2048, pkcs=8, prefix: str = None):
        """
        生成RSA密钥
        @return: 返回生成的密钥内容
        """
        dir = dir if dir.endswith("/") else dir + "/"
        dir = dir + f"{int(time.time())}/" if prefix is None else dir + f"{prefix}-{int(time.time())}/"
        try:
            os.makedirs(dir, exist_ok=True)
            key = RSA.generate(key_size)
            private_key = key.export_key(format="PEM", pkcs=pkcs)
            public_key = key.publickey().export_key(format="PEM")

            with open(f"{dir}/private_pkcs{pkcs}.pem", "wb") as f:
                f.write(private_key)
            with open(f"{dir}/public.pem", "wb") as f:
                f.write(public_key)

            return {
                "private_key": private_key,
                "public_key": public_key
            }
        except Exception as e:
            raise RuntimeError(f"生成RSA密钥失败: {e}")

if __name__ == '__main__':
    pripempath = "/Users/lizhankang/Documents/shouqianba/pems/ka/beta/28lp467052302/client/priKey.pem"
    data = {
        "name": "李四",
        "age": 10,
        "adult": True
    }
    signature = RSAK.signature(data, RSAK.rsakeyfromfile(pripempath))
    print(signature)

    s = "qcJne9O92AJ3hiJFGvZtHKt0CjizpbSETXDpNmIUlSvNdsuAnCuQyJwZB8Ws+AdXgLwq2VMTEX272OMcKP7eYUS5+JAF6dO5270E1Fq9gkewtLqOIRFRmH9k0KGNNKW5gkt7dw0Xd8bqD0n0B04nqTIjfJ9k/bCQ0hwRRDBzRwFHxn+M2yXuFd+YmOUWaLHBHKWw+IPfYD/Q5T7BMYEy2aH6lq6RhP0ZE1gSHtz32GrCO3qCnJ4rXiKT8fJtHVOrtVjkgKyqOW5OhHGU9xB6J9kLup8bX5eFFoY2/RNGGaouM+m6QyzZHkJHpiDMQIHI7LVvZhIvFsE1zRIkaGu9Eg=="
    pubpempath = "/Users/lizhankang/Documents/shouqianba/pems/ka/beta/28lp467052302/client/pubKey.pem"
    result = RSAK.verify(data, s, RSAK.rsakeyfromfile(pubpempath))
    print(result)
