import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad
from Crypto.Random import get_random_bytes
import os
import base64

class VideoEncryptorGUI:
    def __init__(self, master):
        self.master = master
        master.title("视频加密工具")
        master.geometry("500x400")
        
        self.input_file_path = None
        
        # 创建GUI组件
        self.label_title = tk.Label(master, text="视频文件加密工具", font=("Arial", 14, "bold"))
        self.label_title.pack(pady=10)
        
        self.btn_select = tk.Button(master, text="选择视频文件", command=self.select_file, width=20, height=2)
        self.btn_select.pack(pady=5)
        
        self.label_file = tk.Label(master, text="未选择文件", wraplength=400)
        self.label_file.pack(pady=5)
        
        # Key输入
        self.label_key = tk.Label(master, text="AES密钥 (Base64编码或16/24/32字节字符串):")
        self.label_key.pack(pady=(10,0))
        self.entry_key = tk.Entry(master, width=50, show="*")
        self.entry_key.pack(pady=5)
        
        # IV输入
        self.label_iv = tk.Label(master, text="初始化向量IV (Base64编码或16字节字符串):")
        self.label_iv.pack(pady=(10,0))
        self.entry_iv = tk.Entry(master, width=50)
        self.entry_iv.pack(pady=5)
        
        # 生成随机Key和IV按钮
        self.btn_generate = tk.Button(master, text="生成随机密钥和IV", command=self.generate_key_iv, width=20)
        self.btn_generate.pack(pady=5)
        
        # 加密按钮
        self.btn_encrypt = tk.Button(master, text="加密文件", command=self.encrypt_video, width=20, height=2, bg="#4CAF50", fg="white")
        self.btn_encrypt.pack(pady=10)
        
        self.status_label = tk.Label(master, text="", fg="blue")
        self.status_label.pack(pady=5)
        
        # 说明文本
        help_text = "说明：输出文件将与输入文件同目录，扩展名为.hnv"
        self.label_help = tk.Label(master, text=help_text, fg="gray", font=("Arial", 9))
        self.label_help.pack(pady=10)

    def select_file(self):
        """选择要加密的文件"""
        file_path = filedialog.askopenfilename(
            title="选择视频文件",
            filetypes=[("视频文件", "*.mp4 *.avi *.mov *.mkv"), ("所有文件", "*.*")]
        )
        if file_path:
            self.input_file_path = file_path
            self.label_file.config(text=file_path)
            self.status_label.config(text="文件已选择，请设置密钥")

    def generate_key_iv(self):
        """生成随机的密钥和IV"""
        # 生成随机密钥（32字节，AES-256）和IV（16字节）
        key = get_random_bytes(32)
        iv = get_random_bytes(16)
        
        # 转换为Base64编码字符串，显示在输入框中
        key_b64 = base64.b64encode(key).decode('utf-8')
        iv_b64 = base64.b64encode(iv).decode('utf-8')
        
        self.entry_key.delete(0, tk.END)
        self.entry_key.insert(0, key_b64)
        self.entry_iv.delete(0, tk.END)
        self.entry_iv.insert(0, iv_b64)
        
        self.status_label.config(text="已生成随机密钥和IV，请妥善保存！")

    def encrypt_video(self):
        """加密视频文件"""
        if not self.input_file_path:
            messagebox.showerror("错误", "请先选择要加密的文件")
            return
            
        key_str = self.entry_key.get().strip()
        iv_str = self.entry_iv.get().strip()
        
        if not key_str or not iv_str:
            messagebox.showerror("错误", "请输入密钥和IV")
            return
        
        try:
            # 处理密钥：尝试Base64解码，否则使用字符串编码
            try:
                key = base64.b64decode(key_str)
            except:
                key = key_str.encode('utf-8')
            
            # 处理IV：尝试Base64解码，否则使用字符串编码
            try:
                iv = base64.b64decode(iv_str)
            except:
                iv = iv_str.encode('utf-8')
            
            # 确保密钥长度为16、24或32字节
            if len(key) not in [16, 24, 32]:
                if len(key) < 16:
                    key = key.ljust(16, b'\0')
                elif len(key) < 24:
                    key = key.ljust(24, b'\0')
                elif len(key) < 32:
                    key = key.ljust(32, b'\0')
                else:
                    key = key[:32]
            
            # 确保IV长度为16字节
            if len(iv) != 16:
                if len(iv) < 16:
                    iv = iv.ljust(16, b'\0')
                else:
                    iv = iv[:16]
            
            # 生成输出文件路径
            file_dir = os.path.dirname(self.input_file_path)
            file_name = os.path.basename(self.input_file_path)
            file_base = os.path.splitext(file_name)[0]
            output_file = os.path.join(file_dir, f"{file_base}.hnv")
            
            # 执行加密
            self.do_encryption(self.input_file_path, output_file, key, iv)
            
            self.status_label.config(text=f"加密成功！输出文件: {output_file}")
            messagebox.showinfo("成功", f"文件加密完成！\n输出文件: {output_file}")
            
        except Exception as e:
            messagebox.showerror("加密错误", f"加密过程中发生错误: {str(e)}")

    def do_encryption(self, input_path, output_path, key, iv):
        """
        执行加密操作
        使用AES/CBC/PKCS5Padding模式，与Java代码保持一致
        """
        # 读取输入文件
        with open(input_path, 'rb') as f_in:
            plaintext = f_in.read()
        
        # 创建密码器
        cipher = AES.new(key, AES.MODE_CBC, iv)
        
        # 应用PKCS5填充并加密
        padded_data = pad(plaintext, AES.block_size)
        ciphertext = cipher.encrypt(padded_data)
        
        # 写入输出文件
        with open(output_path, 'wb') as f_out:
            f_out.write(ciphertext)

def main():
    # 检查所需库是否已安装
    try:
        from Crypto.Cipher import AES
        from Crypto.Util.Padding import pad
        from Crypto.Random import get_random_bytes
    except ImportError:
        print("正在安装所需库...")
        import subprocess
        import sys
        subprocess.check_call([sys.executable, "-m", "pip", "install", "pycryptodome"])
        print("库安装完成，请重新运行程序")
        return
    
    # 创建主窗口
    root = tk.Tk()
    app = VideoEncryptorGUI(root)
    root.mainloop()

if __name__ == "__main__":
    main()