python版本使用椭圆曲线执行密钥交换

水一篇，

BirdTalk服务端基本快写完了，开始写一个完整的客户端测试；

决定从python入手，因为与其他功能对接时候或者写机器人客服，脚本用的比较多；

直接上代码，原理参考之前的文档。

from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend
import os
import struct
import base64
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

class ECDHKeyExchange:
    def __init__(self):
        self.private_key = None
        self.public_key = None
        self.shared_key = None
        self.key_print = 0

    def generate_key_pair(self):
        """Generate an ECDH key pair."""
        self.private_key = ec.generate_private_key(ec.SECP256R1(), default_backend())
        self.public_key = self.private_key.public_key()
        print("Key pair generated.")

    def export_public_key(self, filename):
        """Export the generated public key to a file."""
        if self.public_key is None:
            raise ValueError("Public key not generated yet.")
        pem = self.public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        if filename != "":
            with open(filename, 'wb') as f:
                f.write(pem)
            print(f"Public key saved to {filename}.")
        return pem.decode('utf-8')
    
    def get_public_key(self):
        if self.public_key is None:
            raise ValueError("Public key not generated yet.")
        pem = self.public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        return pem


    def exchange_keys_from_file(self, peer_public_key_filename):
        """Exchange keys and generate the shared key using the peer's public key."""
        if self.private_key is None:
            raise ValueError("Private key not generated yet.")
        with open(peer_public_key_filename, 'rb') as f:
            peer_public_key_pem = f.read()
        peer_public_key = serialization.load_pem_public_key(peer_public_key_pem, backend=default_backend())
        self.shared_key = self.private_key.exchange(ec.ECDH(), peer_public_key)
        print("Shared key generated.")
    
    def exchange_keys(self, peer_public_key_pem):
        """Exchange keys and generate the shared key using the peer's public key PEM string."""
        if self.private_key is None:
            raise ValueError("Private key not generated yet.")
        peer_public_key = serialization.load_pem_public_key(peer_public_key_pem.encode('utf-8'), backend=default_backend())
        self.shared_key = self.private_key.exchange(ec.ECDH(), peer_public_key)
        print("Shared key generated.")

    def get_shared_key(self):
        if self.shared_key is None:
            raise ValueError("Shared key not generated yet.")
        
        return self.shared_key

    def save_shared_key(self, filename):
        """Save the shared key to a file."""
        if self.shared_key is None:
            raise ValueError("Shared key not generated yet.")
        with open(filename, 'wb') as f:
            f.write(self.shared_key)
        print(f"Shared key saved to {filename}.")
    
    def export_shared_key_base64(self):
        """Export the shared key as a base64 encoded string."""
        if self.shared_key is None:
            raise ValueError("Shared key not generated yet.")
        return base64.b64encode(self.shared_key).decode('utf-8')

    def load_shared_key(self, filename):
        """Load the shared key from a file."""
        with open(filename, 'rb') as f:
            self.shared_key = f.read()
        print(f"Shared key loaded from {filename}.")

    def save_key_print(self, filename):
        """Save the shared key to a file."""
        if self.key_print is None:
            raise ValueError("Shared key not generated yet.")
        with open(filename, 'wb') as f:
            data = str(self.key_print).encode('utf-8')
            f.write(data)
        print(f"key print saved to {filename}.")

    def load_key_print(self, filename):
        """Load the shared key from a file."""
        with open(filename, 'rb') as f:
            data = f.read()
            str = data.decode('utf-8')
            self.key_print = int(str)
        print(f"key print loaded from {filename}.")
        return self.key_print

    def delete_key_file(self, filename):
        """Delete a key file."""
        if os.path.exists(filename):
            os.remove(filename)
            print(f"File {filename} deleted.")
        else:
            print(f"File {filename} does not exist.")
    
    def get_int64_print(self):
        if self.shared_key is None:
            raise ValueError("Shared key not generated yet.")
         
        """Convert a byte array to a 64-bit integer."""
        # 检查字节数组长度是否足够
        if len(self.shared_key) < 8:
            raise ValueError("Insufficient bytes to convert to int64")

        # 将字节数组转换为 int64
        self.key_print = struct.unpack('<q',  self.shared_key[:8])[0]  # 使用 little-endian 格式

        return self.key_print
    
    def encrypt_aes_ctr(self, plaintext):
        if self.shared_key is None:
            raise ValueError("Shared key not generated yet.")
        
        # 生成随机的初始化向量（IV）
        iv = os.urandom(16)  # 初始化向量长度为 16 字节

        # 创建 AES-CTR 算法对象
        algorithm = algorithms.AES(self.shared_key)
        mode = modes.CTR(iv)
        cipher = Cipher(algorithm, mode, backend=default_backend())

        # 使用加密器加密数据
        encryptor = cipher.encryptor()
        encrypted_data = encryptor.update(plaintext) + encryptor.finalize()

        # 将随机初始化向量和加密后的数据拼接在一起
        ciphertext = iv + encrypted_data

        return ciphertext
    
    def decrypt_aes_ctr(self, ciphertext):
        """Decrypt ciphertext using AES-CTR with the shared key."""
        if self.shared_key is None:
            raise ValueError("Shared key not generated yet.")
        
        # 从密文中提取初始化向量（IV）
        iv = ciphertext[:16]  # 初始化向量长度为 16 字节
        encrypted_data = ciphertext[16:]

        # 创建 AES-CTR 算法对象
        algorithm = algorithms.AES(self.shared_key)
        mode = modes.CTR(iv)
        cipher = Cipher(algorithm, mode, backend=default_backend())

        # 使用解密器解密数据
        decryptor = cipher.decryptor()
        plaintext = decryptor.update(encrypted_data) + decryptor.finalize()

        return plaintext
###############################################################
def test_create():
    # 实例化两个 ECDHKeyExchange 对象，模拟两个参与方
    alice = ECDHKeyExchange()
    bob = ECDHKeyExchange()

    # Alice 生成密钥对并导出公钥
    alice.generate_key_pair()
    alice_pub_key = alice.export_public_key("alice_public_key.pem")
    print(alice_pub_key)

    # Bob 生成密钥对并导出公钥
    bob.generate_key_pair()
    bob_pub_key = bob.export_public_key("bob_public_key.pem")
    print(bob_pub_key)

    # Alice 和 Bob 交换公钥并生成共享密钥
    alice.exchange_keys(bob_pub_key)
    bob.exchange_keys(alice_pub_key)

    # 保存共享密钥
    alice.save_shared_key("alice_shared_key.bin")
    bob.save_shared_key("bob_shared_key.bin")

    # 加载共享密钥
    alice.load_shared_key("alice_shared_key.bin")
    bob.load_shared_key("bob_shared_key.bin")
    print(alice.get_int64_print())
    print(bob.get_int64_print())

    alice.save_key_print("alice_key_print.txt")
    keyprint = alice.load_key_print("alice_key_print.txt")
    print(keyprint)

    # 删除密钥文件
    # alice.delete_key_file("alice_public_key.pem")
    # alice.delete_key_file("alice_shared_key.bin")
    # bob.delete_key_file("bob_public_key.pem")
    # bob.delete_key_file("bob_shared_key.bin")

def test_load():
    alice = ECDHKeyExchange()
    alice.load_shared_key("alice_shared_key.bin")
    #print(alice.get_shared_key())
    key_print = alice.load_key_print("alice_key_print.txt")
    print(key_print)

    tm = '123456789412'
    ciper = alice.encrypt_aes_ctr(tm.encode('utf-8'))
    plain = alice.decrypt_aes_ctr(ciper)
    print(plain)

# 示例使用
if __name__ == "__main__":
    test_load()