import socket
from gmssl import sm2
import os


# 收文件


def enc_send(message, sm2_crypt, s):
    enc_message = sm2_crypt.encrypt(message.encode())
    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.bind(('127.0.0.1', 8001))
s.listen()
conn, address = s.accept()
his_public_key = conn.recv(512).decode()
conn.sendall(public_key.encode())
his_sm2_crypt = sm2.CryptSM2(public_key=his_public_key, private_key='')
file_name = dec_recv(my_sm2_crypt, conn, 1024).decode()
enc_send('成功接收文件名' + file_name, his_sm2_crypt, conn)

file_len = int(dec_recv(my_sm2_crypt, conn, 1024).decode())
enc_send('成功接收文件大小' + str(file_len), his_sm2_crypt, conn)

message = dec_recv(my_sm2_crypt, conn, file_len * 2)
enc_send('成功接收文件内容' + file_name, his_sm2_crypt, conn)

data = dec_recv(my_sm2_crypt, conn, 1024).decode()
print('收到对方消息：' + data)
enc_send('已收到文件：' + file_name, his_sm2_crypt, conn)
new_file_name = file_name
while os.path.exists(new_file_name):
    new_file_name = new_file_name[0:new_file_name.rindex('.')] + '(1)' + new_file_name[new_file_name.rindex('.'):]
f = open(new_file_name, 'wb')
f.write(message)
f.close()
print('已收到' + file_name + '保存为' + new_file_name)
s.close()