import socket
import json
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import serialization, hashes


def generate_key():
    return Fernet.generate_key()


def encrypt_file(file_path, key):
    with open(file_path, 'rb') as file:
        data = file.read()
    fernet = Fernet(key)
    encrypted_data = fernet.encrypt(data)
    return encrypted_data


def encrypt_key(public_key, symmetric_key):
    encrypted_key = public_key.encrypt(
        symmetric_key,
        padding.OAEP(
            mgf=padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )
    return encrypted_key


def main():
    host = '127.0.0.1'
    port = 49670

    # 创建rsa密钥对
    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    public_key = private_key.public_key()

    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind((host, port))
    server_socket.listen(1)
    print("Server is listening...")

    conn, addr = server_socket.accept()
    print(f"Connection from {addr}")

    # 发送公钥
    pem = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )
    conn.sendall(pem)

    # 加密文件，这里假设要发送的文件名为 'example.txt'，可根据实际情况修改
    file_path = 'example.txt'
    symmetric_key = generate_key()
    encrypted_data = encrypt_file(file_path, symmetric_key)

    # 用公钥加密对称密钥
    encrypted_key = encrypt_key(public_key, symmetric_key)

    # 构建数字信封，确保数据编码使用更通用的utf-8格式
    envelope = {
        'algorithm': 'AES',
        'key': encrypted_key.hex(),
        'file_data': encrypted_data.decode('utf-8')
    }

    # 发送数字信封，先转换为JSON字符串再编码为字节类型发送
    conn.sendall(json.dumps(envelope).encode('utf-8'))
    conn.close()
    print("File sent successfully.")


if __name__ == "__main__":
    main()