import socket
from gmssl import sm2


# 发文件


def enc_send(message, sm2_crypt, s):
    if type(message) != type(b'ab'):
        message = message.encode()
    enc_message = sm2_crypt.encrypt(message)
    s.sendall(enc_message)


def dec_recv(sm2_crypt, s, l):
    data = s.recv(l)
    dec_data = sm2_crypt.decrypt(data)
    return dec_data


private_key = '00B9AB0B828FF68872F21A837FC303668428DEA11DCD1B24429D0C99E24EED83D5'
public_key = 'B9C9A6E04E9C91F7BA880429273747D7EF5DDEB0BB2FF6317EB00BEF331A83081A6994B8993F3F5D6EADDDB81872266C87C018FB4162F5AF347B483E24620207'
my_sm2_crypt = sm2.CryptSM2(public_key=public_key, private_key=private_key)
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect(('127.0.0.1', 8001))
while True:
    file_name = input('请输入要打开的文件：')
    try:
        f = open(file_name, 'rb')
    except:
        print('找不到' + file_name)
        continue
    break
message = f.read()
f.close()
file_len = str(len(message))
s.sendall(public_key.encode())
his_public_key = s.recv(512).decode()
his_sm2_crypt = sm2.CryptSM2(public_key=his_public_key, private_key='')

enc_send(file_name, his_sm2_crypt, s)

data = dec_recv(my_sm2_crypt, s, 1024).decode()
print('收到对方消息：' + data)

enc_send(file_len, his_sm2_crypt, s)

data = dec_recv(my_sm2_crypt, s, 1024).decode()
print('收到对方消息：' + data)

enc_send(message, his_sm2_crypt, s)

data = dec_recv(my_sm2_crypt, s, 1024).decode()
print('收到对方消息：' + data)

enc_send('发送完毕', his_sm2_crypt, s)

data = dec_recv(my_sm2_crypt, s, 1024).decode()
print('收到对方消息：' + data)
s.close()