import socket
import zipfile
import os

def send_zip(HOST,PORT):
    # 获取当前文件的绝对路径
    current_file_path = os.path.abspath(__file__)

    # 获取当前文件的目录路径
    current_directory = os.path.dirname(current_file_path)

    # 使用相对路径构建其他文件的路径
    file_path = os.path.join(current_directory, 'file.txt.encrypted')
    encrypted_key_path = os.path.join(current_directory, 'encrypted_symmetric_key.bin')
    Apub_key_path = os.path.join(current_directory, 'A_public.txt')
    signature_path = os.path.join(current_directory, 'signature.txt')
    output_zip_path = os.path.join(current_directory, 'encrypted_data.zip')

    # 创建一个 ZIP 文件
    with zipfile.ZipFile(output_zip_path, 'w') as zip_file:
        # 将文件.txt添加到ZIP文件中
        zip_file.write(file_path, arcname='file.txt.encrypted')

        # 将A的公钥添加到ZIP文件中
        zip_file.write(Apub_key_path, arcname='A_public.txt')

        # 将加密后的对称密钥添加到ZIP文件中
        zip_file.write(encrypted_key_path, arcname='encrypted_symmetric_key.bin')

        # 将数字签名添加到ZIP文件中
        zip_file.write(signature_path, arcname='signature.txt')

    print("打包封装完成：", output_zip_path)

    # 设置服务器的IP地址和端口号
    # HOST = host
    # PORT = port

    # 客户端连接服务器
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as client_socket:
        client_socket.connect((HOST, PORT))
        
        # 读取要传输的文件
        with open(output_zip_path, 'rb') as file:
            file_data = file.read()

        # 发送文件数据
        client_socket.sendall(file_data)
        print("文件发送完成")

def receive_Bpubkey(HOST,PORT):
    # 设置服务器的IP地址和端口号
    # HOST = '10.34.52.156'
    # PORT = 22222

    # 创建一个socket对象
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
        # 绑定IP和端口
        server_socket.bind((HOST, PORT))
        # 开始监听传入的连接
        server_socket.listen()

        print(f"等待客户端连接...")
        conn, addr = server_socket.accept()  # 接受客户端的连接
        print(f"已连接：{addr}")

        with conn:
            # 从客户端接收数据
            file_data = conn.recv(1024)  # 假设每次接收的数据大小为1024字节
            with open('B_public.txt', 'wb') as file:
                while file_data:
                    file.write(file_data)
                    file_data = conn.recv(1024)
            print("文件接收完成")

# receive_Bpubkey()

# send_zip()