import gradio as gr
from cryptography.fernet import Fernet
import os

# AES 加密函数（Fernet 实现）
def encrypt_file_gr(file):
    file_path = file.name
    key = Fernet.generate_key()
    cipher = Fernet(key)
    with open(file_path, 'rb') as f:
        data = f.read()
    encrypted_data = cipher.encrypt(data)
    enc_path = file_path + '.enc'
    with open(enc_path, 'wb') as f:
        f.write(encrypted_data)
    return key.decode(), enc_path

# AES 解密函数，恢复原始文件名和格式
def decrypt_file_gr(file, key):
    file_path = file.name
    cipher = Fernet(key.encode())
    try:
        with open(file_path, 'rb') as f:
            encrypted_data = f.read()
        decrypted_data = cipher.decrypt(encrypted_data)
        # 去掉 .enc 后缀，恢复原始文件名
        if file_path.lower().endswith('.enc'):
            dec_path = file_path[:-4]
        else:
            dec_path = file_path + '.dec'
        with open(dec_path, 'wb') as f:
            f.write(decrypted_data)
        return dec_path
    except Exception:
        return None

# Gradio 界面设计
with gr.Blocks() as demo:
    gr.Markdown("# 文件加密解密工具")

    with gr.Tab("加密"):
        encrypt_in = gr.File(label="上传文件")
        encrypt_btn = gr.Button("加密文件 🔒")
        # 文本框显示密钥并提供复制按钮
        encrypt_key = gr.Textbox(label="生成的密钥", interactive=False, show_copy_button=True)
        encrypt_out = gr.File(label="下载加密文件 (.enc)")
        encrypt_btn.click(
            fn=encrypt_file_gr,
            inputs=encrypt_in,
            outputs=[encrypt_key, encrypt_out]
        )

    with gr.Tab("解密"):
        decrypt_in = gr.File(label="上传加密文件 (.enc)")
        decrypt_key_in = gr.Textbox(label="输入密钥")
        decrypt_btn = gr.Button("解密文件 🔓")
        decrypt_out = gr.File(label="下载解密文件（原始格式）")
        decrypt_btn.click(
            fn=decrypt_file_gr,
            inputs=[decrypt_in, decrypt_key_in],
            outputs=decrypt_out
        )

if __name__ == "__main__":
    demo.launch()