from Crypto.PublicKey import RSA
from Crypto.Cipher import AES, PKCS1_OAEP, DES3
from Crypto.Hash import SHA256
from Crypto.Signature import pkcs1_15
import base64
import json
import tkinter as tk
from tkinter import filedialog, messagebox
from tkinter.ttk import Combobox

def decrypt_file(digital_envelope, private_key_path, symmetric_algorithm='AES', mode='CBC'):
    """
    使用包含加密对称密钥和加密文件数据的数字信封解密文件，并验证文件完整性。
    """
    try:
        with open(private_key_path, 'r') as f:
            private_key = RSA.import_key(f.read())

        cipher_rsa = PKCS1_OAEP.new(private_key)
        encrypted_symmetric_key = base64.b64decode(digital_envelope['encrypted_symmetric_key'])
        symmetric_key = cipher_rsa.decrypt(encrypted_symmetric_key)

        if symmetric_algorithm == 'AES':
            if mode == 'ECB':
                cipher_aes = AES.new(symmetric_key, AES.MODE_ECB)
            elif mode == 'CBC':
                iv = base64.b64decode(digital_envelope['iv'])
                cipher_aes = AES.new(symmetric_key, AES.MODE_CBC, iv)
            else:
                raise ValueError("不支持的模式")

            encrypted_file_data = base64.b64decode(digital_envelope['encrypted_file_data'])

            decrypted_padded_file_data = cipher_aes.decrypt(encrypted_file_data)
            if mode == 'ECB' or mode == 'CBC':
                padding_length = decrypted_padded_file_data[-1]
                decrypted_file_data = decrypted_padded_file_data[:-padding_length]

        elif symmetric_algorithm == 'DES3':
            if mode == 'ECB':
                cipher_des3 = DES3.new(symmetric_key, DES3.MODE_ECB)
            elif mode == 'CBC':
                iv = base64.b64decode(digital_envelope['iv'])
                cipher_des3 = DES3.new(symmetric_key, DES3.MODE_CBC, iv)
            else:
                raise ValueError("不支持的模式")

            encrypted_file_data = base64.b64decode(digital_envelope['encrypted_file_data'])

            decrypted_padded_file_data = cipher_des3.decrypt(encrypted_file_data)
            if mode == 'ECB' or mode == 'CBC':
                padding_length = decrypted_padded_file_data[-1]
                decrypted_file_data = decrypted_padded_file_data[:-padding_length]

        else:
            raise ValueError("不支持的算法")

        hash_obj = SHA256.new(decrypted_file_data)
        try:
            pkcs1_15.new(private_key).verify(hash_obj, base64.b64decode(digital_envelope['signature']))
            print("签名验证成功")
        except (ValueError, TypeError):
            print("签名验证失败")
            return None

        return decrypted_file_data

    except Exception as e:
        messagebox.showerror("错误", str(e))
        return None

def select_digital_envelope():
    file_path = filedialog.askopenfilename(title="选择数字信封文件", filetypes=[("JSON files", "*.json")])
    digital_envelope_entry.delete(0, tk.END)
    digital_envelope_entry.insert(0, file_path)

def on_decrypt_button_click():
    file_path = digital_envelope_entry.get()
    if not file_path:
        messagebox.showwarning("警告", "请选择数字信封文件")
        return

    symmetric_algorithm = algorithm_combobox.get()
    mode = mode_combobox.get()

    if not symmetric_algorithm or not mode:
        messagebox.showwarning("警告", "请选择对称加密算法和模式")
        return

    try:
        with open(file_path, 'r') as f:
            digital_envelope = json.load(f)

        decrypted_file_data = decrypt_file(digital_envelope, 'sender_private.pem', symmetric_algorithm, mode)

        if decrypted_file_data is not None:
            save_path = filedialog.asksaveasfilename(defaultextension=".txt", filetypes=[("Text files", "*.txt")])
            if save_path:
                with open(save_path, 'wb') as f:
                    f.write(decrypted_file_data)
                messagebox.showinfo("成功", "文件已成功解密并保存")
        else:
            messagebox.showerror("错误", "解密失败或签名验证未通过")

    except Exception as e:
        messagebox.showerror("错误", str(e))

def main():
    global digital_envelope_entry, algorithm_combobox, mode_combobox

    root = tk.Tk()
    root.title("文件解密工具")

    # 数字信封文件选择部分
    digital_envelope_label = tk.Label(root, text="选择数字信封文件:")
    digital_envelope_label.grid(row=0, column=0, padx=10, pady=10, sticky=tk.W)

    digital_envelope_entry = tk.Entry(root, width=50)
    digital_envelope_entry.grid(row=0, column=1, padx=10, pady=10)

    browse_button = tk.Button(root, text="浏览", command=select_digital_envelope)
    browse_button.grid(row=0, column=2, padx=10, pady=10)

    # 对称加密算法选择部分
    algorithm_label = tk.Label(root, text="选择对称加密算法:")
    algorithm_label.grid(row=1, column=0, padx=10, pady=10, sticky=tk.W)

    algorithm_combobox = Combobox(root, values=['AES', 'DES3'], state='readonly')
    algorithm_combobox.set('AES')
    algorithm_combobox.grid(row=1, column=1, padx=10, pady=10)

    # 加密模式选择部分
    mode_label = tk.Label(root, text="选择加密模式:")
    mode_label.grid(row=2, column=0, padx=10, pady=10, sticky=tk.W)

    mode_combobox = Combobox(root, values=['ECB', 'CBC'], state='readonly')
    mode_combobox.set('CBC')
    mode_combobox.grid(row=2, column=1, padx=10, pady=10)

    # 解密按钮
    decrypt_button = tk.Button(root, text="解密文件", command=on_decrypt_button_click)
    decrypt_button.grid(row=3, columnspan=3, pady=20)

    root.mainloop()

if __name__ == "__main__":
    main()