import easygui
import my_sm9 as sm9
import optimized_field_elements as fq
import pickle

def my_input(msg, title, fieldNames):
    fieldValues = easygui.multenterbox(msg, title, fieldNames)
    while True:
        if fieldValues == None:
            break
        errmsg = ""
        for i in range(len(fieldNames)):
            option = fieldNames[i].strip()
            if fieldValues[i].strip() == "" and option[0] == "*":
                errmsg += ("【%s】为必填项   " % fieldNames[i])
        if errmsg == "":
            break
        fieldValues = easygui.multenterbox(errmsg, title, fieldNames, fieldValues)
    return fieldValues


def read_pub():
    msg = "请填写以下信息"
    title = '读取公钥文件'
    fieldNames = ["*公钥文件路径"]
    fieldValues = my_input(msg, title, fieldNames)
    if fieldValues == None:
        easygui.msgbox(msg='不能为空！', title='错误')
        return False
    path = fieldValues[0]
    try:
        file = open(path, 'rb')
    except:
        easygui.msgbox(msg='找不到文件！', title='错误')
        return False
    try:
        user, master_public = pickle.load(file)
    except:
        easygui.msgbox(msg='文件格式错误！', title='错误')
        return False
    file.close()
    easygui.msgbox(msg='公私钥文件%spub已读取' % user, title='结果')
    return user, master_public


user_list=[]
master_public_list=[]
my_user=my_master_public=my_master_secret=None
try:
    file = open('my_sec.txt', 'rb')
    my_user,my_master_public,my_master_secret = pickle.load(file)
    file.close()
except:
    pass
print(my_user)
try:
    file=open('user_list.txt','rb')
    user_list,master_public_list = pickle.load(file)
    file.close()
except:
    pass
print(user_list)
print(master_public_list)
while True:
    a=''
    if my_user:
        a+=my_user+'，'
    choice = easygui.buttonbox(msg=a+'欢迎使用，请选择要使用的功能', title='主菜单', choices=('生成用户密钥文件', '读取公钥文件', '读取公私钥文件', '加解密', '签名验签', '退出程序'))
    if choice == '生成用户密钥文件':
        msg = "请填写一下信息(其中带*号的项为必填项)"
        title = '生成用户密钥文件'
        fieldNames = ["*用户标识"]
        fieldValues = my_input(msg, title, fieldNames)
        if fieldValues == None:
            continue
        my_user=fieldValues[0]
        p1 = fq.G2
        p2 = fq.G1
        s_public, s_g, s_secret = sm9.sign_setup(p1, p2)
        e_public, e_g, e_secret = sm9.encrypt_setup(p1, p2)
        my_master_public = (p1, p2, s_public, s_g, e_public, e_g)
        my_master_secret = (s_secret, e_secret)
        file_pub = open('my_pub.txt', 'wb')
        file_sec = open('my_sec.txt', 'wb')
        pickle.dump((my_user, my_master_public, my_master_secret), file_sec)
        pickle.dump((my_user, my_master_public), file_pub)
        easygui.msgbox(msg='用户文件%s、%s已生产'%('my_pub','my_sec'), title='结果')
        file_pub.close()
        file_sec.close()
    elif choice == '读取公钥文件':
        i=read_pub()
        if not i:
            continue
        user, master_public = i
        if user in user_list:
            i = user_list.index(user)
            user_list[i] = user
            master_public_list[i] = master_public
        else:
            user_list.append(user)
            master_public_list.append(master_public)
        file = open('user_list.txt', 'wb')
        pickle.dump((user_list,master_public_list), file)
        file.close()
    elif choice == '读取公私钥文件':
        msg = "请填写一下信息(其中带*号的项为必填项)"
        title = '读取公私钥文件'
        fieldNames = ["*公私钥文件路径"]
        fieldValues = my_input(msg, title, fieldNames)
        if fieldValues == None:
            continue
        path = fieldValues[0]
        file = open(path, 'rb')
        my_user, my_master_public, my_master_secret = pickle.load(file)
        easygui.msgbox(msg='公私钥文件my_sec已读取', title='结果')
        file_pub = open('my_pub.txt', 'wb')
        file_sec = open('my_sec.txt', 'wb')
        pickle.dump((my_user, my_master_public, my_master_secret), file_sec)
        pickle.dump((my_user, my_master_public), file_pub)
        file_pub.close()
        file_sec.close()
        file.close()
    elif my_user or master_public_list!=[]:
        if choice == '加解密':
            e_da= None
            while True:
                en_choice = easygui.buttonbox(msg='请选择要使用的功能', title='加解密菜单',
                                               choices=('用户私钥生成', '加密', '解密', '返回上一级菜单'))
                if en_choice == '用户私钥生成':
                    if not my_master_secret:
                        easygui.msgbox(msg='请先生成用户密钥文件', title='错误')
                        break
                    e_da = sm9.secret_key_extract(my_master_public[1], my_master_secret[1], my_user)  # 计算dsA：用户的加密私钥
                    msg = '生成私钥成功'
                    easygui.msgbox(msg=msg, title='密钥生成结果')
                elif en_choice == '加密':
                    m = easygui.enterbox(msg='请输入明文', title='加密')
                    if not m:
                        continue
                    choose_list=user_list+['添加']
                    usr_choice=easygui.multchoicebox(msg="请选择接收方",title="选择加密对象",choices=choose_list)[0]
                    if usr_choice=='添加':
                        i = read_pub()
                        if not i:
                            continue
                        user, master_public = i
                        if user in user_list:
                            i=user_list.index(user)
                            user_list[i]=user
                            master_public_list[i]=master_public
                        else:
                            user_list.append(user)
                            master_public_list.append(master_public)
                        file = open('user_list.txt', 'wb')
                        pickle.dump((user_list, master_public_list), file)
                        file.close()
                        this_user=user
                        this_master_public=master_public
                    elif not usr_choice:
                        continue
                    else:
                        this_user=usr_choice
                        this_master_public=master_public_list[user_list.index(this_user)]
                    c=sm9.kem_dem_enc(this_master_public, this_user, m.encode(), 32)
                    file_mm=open('mw.txt','wb')
                    pickle.dump((my_user,c), file_mm)
                    file_mm.close()
                    easygui.msgbox(msg='加密结果保存于mw.txt中' , title='加密结果')
                elif en_choice == '解密':
                    if not e_da:
                        easygui.msgbox(msg='请先生成解密私钥', title='错误')
                        continue
                    path_mw = easygui.enterbox(msg='请输入密文的路径', title='解密')
                    if not path_mw:
                        continue
                    file_mm=open(path_mw,'rb')
                    this_user,c=pickle.load(file_mm)
                    file_mm.close()
                    try:
                        m=sm9.kem_dem_dec(my_user, e_da, c, 32).decode()
                    except:
                        easygui.msgbox(msg='解密失败', title='错误')
                        continue
                    easygui.msgbox(msg=path_mw + '的解密结果为' + m+'。发送方为'+this_user, title='解密结果')
                else:
                    break
        elif choice == '签名验签':
            s_da = None
            while True:
                sign_choice = easygui.buttonbox(msg='请选择要使用的功能', title='签名验签菜单',
                                               choices=('私钥生成', '签名', '验签', '返回上一级菜单'))
                if sign_choice == '私钥生成':
                    s_da = sm9.secret_key_extract(my_master_public[0], my_master_secret[0], my_user)  # 计算dsA：用户 A 的签名私钥
                    msg = '生成私钥成功'
                    easygui.msgbox(msg=msg, title='密钥生成结果')
                elif sign_choice == '签名':
                    m = easygui.enterbox(msg='请输入明文', title='签名')
                    if not m:
                        continue
                    signature = sm9.sign(my_master_public[3], s_da, m.encode())  # 签名
                    file_si = open('sign.txt', 'wb')
                    pickle.dump((m,my_user, signature), file_si)
                    file_si.close()
                    easygui.msgbox(msg='加密结果保存sign.txt中' , title='签名结果')
                elif sign_choice == '验签':
                    path_mw = easygui.enterbox(msg='请输入签名的路径', title='解密')
                    if not path_mw:
                        continue
                    file_si = open(path_mw, 'rb')
                    m,this_user,signature = pickle.load(file_si)
                    if this_user in user_list:
                        this_master_public=master_public_list[user_list.index(this_user)]
                    else:
                        easygui.msgbox(msg='没有找到该签名的公钥，请添加', title='错误')
                        i = read_pub()
                        if not i:
                            continue
                        user, master_public=i
                        if user!=this_user:
                            easygui.msgbox(msg='输入的公钥不是该签名的公钥', title='错误')
                            continue
                        user_list.append(user)
                        master_public_list.append(master_public)
                        file = open('user_list.txt', 'wb')
                        pickle.dump((user_list, master_public_list), file)
                        file.close()
                        this_user = user
                        this_master_public = master_public
                    file_si.close()
                    result=sm9.verify(this_master_public, this_user, m.encode(), signature)
                    easygui.msgbox(msg='验签结果为' + str(result), title='验签结果')
                else:
                    break
        elif choice in ('退出程序', None):
            break
    elif choice in('退出程序' ,None):
            break
    else:
        easygui.msgbox(msg='请先生成或读取密钥！', title='错误')