import base64
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.backends import default_backend


def encrypt_message_with_public_key(public_key_pem: str, message: str) -> str:
    """
    用公钥（字符串形式）对明文进行加密，并返回 base64 编码后的密文字符串
    """
    # 1. 从公钥字符串还原为公钥对象
    public_key = serialization.load_pem_public_key(
        public_key_pem.encode("utf-8"),  # 字符串转 bytes
        backend=default_backend()
    )
    
    # 2. 对明文进行加密（OAEP 填充 + SHA256），返回二进制密文
    ciphertext_bytes = public_key.encrypt(
        message.encode("utf-8"),
        padding.PKCS1v15()
    )
    
    # 3. 将二进制密文转成 base64 字符串，方便通过 HTTP 或其他渠道传输
    ciphertext_base64 = base64.b64encode(ciphertext_bytes).decode("utf-8")
    
    return ciphertext_base64


def decrypt_message_with_private_key(private_key_pem: str, ciphertext_base64: str) -> str:
    """
    用私钥（字符串形式）对 base64格式的密文进行解密，返回明文字符串
    """
    # 1. 从私钥字符串还原为私钥对象
    private_key = serialization.load_pem_private_key(
        private_key_pem.encode("utf-8"),  # 字符串转 bytes
        password=None,  # 如果在生成私钥时使用了加密算法，这里需要提供密码
        backend=default_backend()
    )
    
    # 2. Base64解码出二进制密文
    ciphertext_bytes = base64.b64decode(ciphertext_base64)
    
    # 3. 使用私钥解密
    plaintext_bytes = private_key.decrypt(
        ciphertext_bytes,
        padding.PKCS1v15()
    )
    
    return plaintext_bytes.decode("utf-8")