from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend
from loguru import logger


def generate_and_save_keys(private_key_filename: str = "private_key.pem",
                           public_key_filename: str = "public_key.pem",
                           password: str = None):
    """生成RSA公鑰和私鑰對(duì)，并將它們保存到本地文件，私鑰可以選擇性地使用密碼加密。
    :param private_key_filename: 私鑰保存的文件名。
    :param public_key_filename: 公鑰保存的文件名。
    :param password: 用于加密私鑰的密碼。如果為None，私鑰將不加密。
    """
    # 1. 生成密鑰對(duì)
    private_key = rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,  # 推薦使用2048位或更高
        backend=default_backend()
    )
    public_key = private_key.public_key()

    # 2. 私鑰加密
    if password:
        encryption_algorithm = serialization.BestAvailableEncryption(password.encode('utf-8'))
    else:
        encryption_algorithm = serialization.NoEncryption()

    # 生成私鑰
    private_pem = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=encryption_algorithm
    )
    with open(private_key_filename, "wb") as f:
        f.write(private_pem)
    logger.debug(f"私鑰已保存到：{private_key_filename}")

    # 生成公鑰
    public_pem = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )
    with open(public_key_filename, "wb") as f:
        f.write(public_pem)
    logger.debug(f"公鑰已保存到：{public_key_filename}")


if __name__ == "__main__":
    generate_and_save_keys("my_private_key.pem", "my_public_key.pem", password="my_secret_password")

import os
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives import serialization
from loguru import logger
from cryptography.hazmat.backends import default_backend


def generate_random_message(length=100):
    """隨機(jī)生成一段信息"""
    return ''.join(os.urandom(1).hex() for _ in range(length // 2))


def load_public_key(filename):
    """加載公鑰"""
    with open(filename, "rb") as f:
        pem_data = f.read()
    return serialization.load_pem_public_key(pem_data, backend=default_backend())


def encrypt_data(public_key, message):
    """使用公鑰加密消息"""
    encoded_message = message.encode('utf-8')
    ciphertext = public_key.encrypt(
        encoded_message,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )
    return ciphertext


def main():
    send_message = generate_random_message(100)  # 要發(fā)送的信息
    logger.debug(f"明文：{send_message}")
    public_key_object = load_public_key("my_public_key.pem")
    encrypted_msg = encrypt_data(public_key_object, send_message)
    logger.debug(f"密文: {encrypted_msg.hex()}")  # 以十六進(jìn)制顯示密文
    # 這里保存密文
    with open("encrypted_message.bin", "wb") as f:
        f.write(encrypted_msg)

if __name__ == "__main__":
    main()

from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from loguru import logger


def load_private_key(filename, password=None):
    """從文件加載私鑰"""
    with open(filename, "rb") as f:
        pem_data = f.read()
    # 密碼需要編碼為字節(jié)
    password_bytes = password.encode('utf-8') if password else None
    return serialization.load_pem_private_key(
        pem_data, password=password_bytes, backend=default_backend()
    )


def decrypt_data(private_key, ciphertext):
    """使用私鑰解密密文"""
    # 使用OAEP填充模式進(jìn)行解密，與加密時(shí)使用的填充模式保持一致
    plaintext = private_key.decrypt(
        ciphertext,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )
    return plaintext.decode('utf-8')


def main():
    private_key_object = load_private_key("my_private_key.pem", password="my_secret_password")  # 加載私鑰
    with open("encrypted_message.bin", "rb") as f:
        encrypted_msg_from_alice = f.read()
    logger.debug(f"密文:{encrypted_msg_from_alice.hex()}")
    decrypted_message = decrypt_data(private_key_object, encrypted_msg_from_alice)
    logger.debug(f"明文:{decrypted_message}")


if __name__ == "__main__":
    main()