import base64
import json
import socket
import sys
import tkinter.filedialog

import EncryptUtil


# 传输文件的入口
def transfer_file():
    my_username = input("请输入你的用户名：")
    target_username = input("请输入对端用户名：")
    # 生成对话框，询问需要传输的文件
    root = tkinter.Tk()
    root.withdraw()
    filename = tkinter.filedialog.askopenfilename()
    # 使用二进制打开文件
    with open(filename, "rb") as fp:
        file_data = fp.read()
    # 计算文件摘要
    hash_md5 = EncryptUtil.cal_md5(file_data)
    # 生成用于加密文件的公钥
    des_key = EncryptUtil.generate_des_key()
    # 使用des加密文件
    encrypt_data = EncryptUtil.des_encrypt(des_key, file_data)
    # 使用对方公钥对des密钥进行加密
    encrypt_des_key = EncryptUtil.rsa_encrypt(des_key, target_username)
    # 对文件摘要进行签名
    file_sign = EncryptUtil.rsa_sign(hash_md5.encode(), my_username)
    # 使用base64编码
    encrypt_data_b64 = base64.b64encode(encrypt_data).decode()
    encrypt_des_key_b64 = base64.b64encode(encrypt_des_key).decode()
    file_sign_b64 = base64.b64encode(file_sign).decode()
    # 使用json打包
    total_data = {
        "encrypt_data_b64": encrypt_data_b64,
        "encrypt_des_key_b64": encrypt_des_key_b64,
        "file_sign_b64": file_sign_b64
    }
    # json转json string
    total_data_json_str = json.dumps(total_data)
    # 发送打包的数据
    client = socket.socket()
    ip = input("请输入要发送的ip: ")
    client.connect((ip, 8787))
    client.sendall(total_data_json_str.encode())
    client.close()


# 接受文件入口
def receive_file():
    my_username = input("请输入你的用户名：")
    target_username = input("请输入对端用户名： ")
    # 建立服务器socket
    server = socket.socket()
    port = 8787
    # 监听8787端口
    server.bind(("0.0.0.0", port))
    server.listen()
    # 接收连接
    client, address = server.accept()
    # 获取打包的json数据
    client_data = client.recv(10240).decode()
    total_data = json.loads(client_data)
    # 解包
    encrypt_data_b64 = total_data["encrypt_data_b64"]
    encrypt_des_key_b64 = total_data["encrypt_des_key_b64"]
    file_sign_b64 = total_data["file_sign_b64"]
    # 解base64
    encrypt_data = base64.b64decode(encrypt_data_b64)
    encrypt_des_key = base64.b64decode(encrypt_des_key_b64)
    file_sign = base64.b64decode(file_sign_b64)
    # 解密获取des密钥
    des_key = EncryptUtil.rsa_decrypt(encrypt_des_key, my_username)
    # 解密文件
    file_data = EncryptUtil.des_decrypt(des_key, encrypt_data)
    # 计算文件散列值
    hash_md5 = EncryptUtil.cal_md5(file_data)
    file_md5 = EncryptUtil.rsa_verify(file_sign, target_username).decode()
    if hash_md5 == file_md5:
        print("签名和哈希验证成功！")
    else:
        print("签名和哈希验证失败，文件可能被篡改！")
        sys.exit()
    # 验证成功后保存文件
    with open("recv_file", "wb") as fp:
        fp.write(file_data)


# 初始化生成密钥
def generate_key():
    # 在本地生成rsa公私钥对
    username = input("请输入用户名：")
    EncryptUtil.generate_keypair(username)
    print("生成成功！")


if __name__ == '__main__':
    print("请选择操作：")
    print("1.生成密钥")
    print("2.接受文件")
    print("3.传输文件")
    mode = input("请选择：")
    if mode == "1":
        generate_key()
    elif mode == "2":
        receive_file()
    elif mode == "3":
        transfer_file()
    else:
        print("仅可以选择1，2，3")