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


class RSAKey:
    """
    RSA 密钥类：封装 RSA 密钥对象，提供面向对象的签名、验签及密钥管理功能。
    """
    def __init__(self, key: RSA.RsaKey):
        self.__key = key

    @classmethod
    def from_pem(cls, file_path: str) -> "RSAKey":
        """从 PEM 文件加载密钥"""
        with open(file_path, mode="r", encoding="utf-8") as f:
            content = f.read().strip()
        return cls(RSA.import_key(content.encode("utf-8")))

    @classmethod
    def from_text(cls, key_text: str) -> "RSAKey":
        """从密钥文本加载密钥"""
        return cls(RSA.import_key(key_text.encode("utf-8")))

    def sign(self, data: dict) -> str:
        """使用当前私钥对数据进行签名"""
        if not self.is_private:
            raise ValueError("当前密钥不包含私钥，无法进行签名操作")
        
        signer = PKCS1_v1_5.new(self.__key)
        data_bytes = json.dumps(data, separators=(',', ':'), ensure_ascii=False).encode('utf-8')
        data_hash = SHA256.new(data_bytes)
        return base64.b64encode(signer.sign(data_hash)).decode('utf-8')

    def verify(self, data: dict, signature: str) -> bool:
        """使用当前公钥验证数据签名"""
        verifier = PKCS1_v1_5.new(self.__key)
        signature_bytes = base64.b64decode(signature.encode('utf-8'))
        data_bytes = json.dumps(data, separators=(',', ':'), ensure_ascii=False).encode('utf-8')
        data_hash = SHA256.new(data_bytes)
        return verifier.verify(data_hash, signature_bytes)

    @property
    def is_private(self) -> bool:
        """判断当前是否为私钥"""
        return self.__key.has_private()


    def publickey(self) -> "RSAKey":
        """获取当前密钥对应的公钥对象"""
        return RSAKey(self.__key.publickey())

    def export(self, format: str = "PEM", pkcs: int = 8) -> str:
        """导出密钥为字符串文本，密钥主体部分将合并为单行"""
        pem = self.__key.export_key(format=format, pkcs=pkcs).decode("utf-8")
        if format.upper() == "PEM":
            lines = pem.strip().splitlines()
            if len(lines) > 2:
                header = lines[0]
                footer = lines[-1]
                body = "".join(lines[1:-1])
                return f"{header}\n{body}\n{footer}"
        return pem

    def save(self, file_path: str, format: str = "PEM", pkcs: int = 8):
        """将密钥保存到文件"""
        directory = os.path.dirname(file_path)
        if directory:
            os.makedirs(directory, exist_ok=True)
        with open(file_path, "wb") as f:
            f.write(self.__key.export_key(format=format, pkcs=pkcs))

    @staticmethod
    def generate(key_size: int = 2048) -> "RSAKey":
        """生成一个新的 RSA 密钥对"""
        return RSAKey(RSA.generate(key_size))

class RsaKeyBuilder:
    """
    RsaKey 密钥构造器：用于灵活配置并生成或加载 RSAKey 对象。
    """
    def __init__(self):
        self._pkcs = 8
        self._key_size = 2048
        self._key_text = None
        self._file_path = None

    def key_size(self, size: int) -> "RsaKeyBuilder":
        self._key_size = size
        return self

    def pkcs(self, pkcs: int) -> "RsaKeyBuilder":
        self._pkcs = pkcs
        return self

    def keytext(self, text: str) -> "RsaKeyBuilder":
        self._key_text = text
        return self

    def filepath(self, path: str) -> "RsaKeyBuilder":
        self._file_path = path
        return self

    def build(self) -> RSAKey:
        """根据配置构建 RSAKey 对象"""
        if self._key_text:
            return RSAKey.from_text(self._key_text)
        if self._file_path:
            return RSAKey.from_pem(self._file_path)
        return RSAKey.generate(self._key_size)
