from gmssl import sm4
import socket

PORT = 8081
BUFFER_SIZE = 1024
EXIT_SYMBOL = b"exit"
key = b"1234567890abcdef"

def encrypt_sm4(key, ptext):
    """使用SM4算法加密数据"""
    crypt_sm4 = sm4.CryptSM4()
    crypt_sm4.set_key(key, sm4.SM4_ENCRYPT)
    ctext = crypt_sm4.crypt_ecb(ptext)
    return ctext

sockfd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
    server_addr = ('127.0.0.1', PORT)
    sockfd.connect(server_addr)
    print("Connected to server!")
    while True:
        message = input("Enter message (type 'exit' to quit): ").encode("utf-8")
        ctext = encrypt_sm4(key, message)
        print(f"加密后的消息: {ctext.hex()}")
        try:
            sockfd.sendall(ctext)
        except socket.error as e:
            print(f"write to server: {e}")
            break
        # 检查是否输入了退出符号
        if message.startswith(EXIT_SYMBOL):
            print("Sending exit command to server.")
            break
        # 从服务器接收数据
        try:
            data = sockfd.recv(BUFFER_SIZE - 1)
            if not data:
                break
            buffer = data.decode()
            buffer = buffer.rstrip('\0')
            print(f"Received from server: {buffer}")
        except socket.error as e:
            print(f"read from server: {e}")
            break
except socket.error as e:
    print(f"Socket error: {e}")
finally:
    # 关闭socket
    sockfd.close()
    print("Client closed.")