import gradio as gr
from cryptography.fernet import Fernet

# AES 加密函数（Fernet 实现）
def encrypt_file_gr(file):
    # `file` 是一个 NamedTemporaryFile 对象，带有属性 `name`
    file_path = file.name
    # 生成密钥并初始化 cipher
    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)
        dec_path = file_path.replace('.enc', '.dec')
        with open(dec_path, 'wb') as f:
            f.write(decrypted_data)
        return dec_path
    except Exception as e:
        # 如果解密失败，返回空，让 Gradio 显示错误
        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)
        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="下载解密文件 (.dec)")
        decrypt_btn.click(
            fn=decrypt_file_gr,
            inputs=[decrypt_in, decrypt_key_in],
            outputs=decrypt_out
        )

# 本地服务器启动函数，Hugging Face Spaces 会自动调用 demo.launch()
if __name__ == "__main__":
    demo.launch()