from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
import os
import socket

# 密钥和初始化向量，需和服务端保持一致
key = b'DKYERDKYERDKYERS'
iv = b'dkyerdkyerdkyers'

def aes_encrypt(data):
    try:
        cipher = AES.new(key, AES.MODE_CBC, iv)
        padded_data = pad(data.encode('utf-8'), AES.block_size)
        encrypted_data = cipher.encrypt(padded_data)
        return encrypted_data
    except Exception as e:
        print(f"加密过程出错: {e}")
        return None

server_address = ('192.168.92.1', 8888)
file_name = 'send_file.txt'
log_file = 'client_log.txt'
client_socket = None

try:
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.connect(server_address)

    script_dir = os.path.dirname(os.path.abspath(__file__))
    file_path = os.path.join(script_dir, file_name)

    # 每次都新建文件并输入内容
    print("\n请输入要发送的内容（输入空行结束）:")
    content = []
    while True:
        line = input()
        if not line:
            break
        content.append(line)
    content = '\n'.join(content)
    with open(file_path, 'w', encoding='utf-8') as file:
        file.write(content)
    print(f"内容已保存到 {file_name}")

    with open(file_path, 'r', encoding='utf-8') as file:
        plaintext_data = file.read()

    encrypted_data = aes_encrypt(plaintext_data)
    if encrypted_data is None:
        raise ValueError("加密失败，无法继续执行。")

    with open(log_file, 'w', encoding='utf-8') as log:
        log.write("加密前的内容:\n")
        log.write(plaintext_data)
        log.write("\n\n加密后的内容:\n")
        log.write(encrypted_data.hex())

    confirm = input("\n是否发送文件内容？(Y/N): ").strip().upper()
    if confirm == 'Y':
        client_socket.send(encrypted_data)
        print('数据已加密并发送')
    else:
        print("发送已取消。")

    # 询问用户是否发送已有文件
    send_existing_file = input("\n是否发送已有本地文件？(Y/N): ").strip().upper()
    if send_existing_file == 'Y':
        existing_file_path = input("请输入要发送的文件路径: ").strip()
        if os.path.isfile(existing_file_path):
            with open(existing_file_path, 'r', encoding='utf-8') as file:
                existing_file_content = file.read()
            existing_file_encrypted_data = aes_encrypt(existing_file_content)
            if existing_file_encrypted_data is not None:
                client_socket.send(existing_file_encrypted_data)
                print(f"文件 {existing_file_path} 的内容已加密并发送。")
            else:
                print("加密已有文件内容失败。")
        else:
            print("文件路径无效或文件不存在。")

except ConnectionRefusedError:
    print("无法连接到服务器，请检查服务端是否启动。")
except FileNotFoundError:
    print("文件未找到，请检查文件路径。")
except Exception as e:
    print(f"发生未知错误: {e}")
finally:
    if client_socket:
        client_socket.close()