import socket
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import base64

Key = b'1234567890123456'
iv = b'1234567890123456'

def encrypt(plain_text, key, iv):
    cipher = AES.new(key, AES.MODE_CBC, iv)
    padded_text = pad(plain_text.encode(), AES.block_size)
    encrypted_text = cipher.encrypt(padded_text)
    return base64.b64encode(encrypted_text).decode()

def decrypt(encrypted_text, key, iv):
    cipher = AES.new(key, AES.MODE_CBC, iv)
    decrypted_text = unpad(cipher.decrypt(base64.b64decode(encrypted_text)), AES.block_size)
    return decrypted_text.decode()

# 创建服务器 socket
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
address = ('localhost', 1234)
s.bind(address)
s.listen(5)
print("服务器启动，等待{}连接...".format(address))

# 接受连接
c, addr = s.accept()
print("我们收到了{}的连接".format(addr))

with open("server_output.txt", "w", encoding="utf-8") as f:
    while True:
        # 接收数据
        data = c.recv(1024).decode("utf-8")
        if not data:
            break

        # 解密数据
        try:
            decrypted_text = decrypt(data, Key, iv)
            print("服务端收到：{}".format(decrypted_text))
            f.write(decrypted_text + "\n")
            f.flush()
        except Exception as e:
            print(f"解密失败：{e}")
            break

        # 获取服务端回应，加密后发送
        response = input("服务端回应：")
        encrypted_response = encrypt(response, Key, iv)
        c.send(encrypted_response.encode("utf-8"))

        # 是否继续
        if input("是否继续？(y/n):") == 'n':
            break

c.close()
s.close()