app.py

import os
import io
import gradio as gr
import json
import tempfile
import zipfile
from PIL import Image

# --- Required Core Logic ---
# These functions are copied from your project's core modules.
# Ensure keylock1 and repo_to_md are in the same directory or accessible in the Python path.
from keylock import core as keylock_core
from repo_to_md.core import markdown_to_files

# --- Helper Functions for the Decoder UI ---

def retrieve_and_process_files(stego_image_pil: Image.Image, password: str):
    """
    Extracts data from an image and prepares the file browser UI components for update.
    """
    # Initial state for all UI components
    status = "An error occurred."
    file_data_state, file_buffers_state = [], {}
    file_browser_visibility = gr.update(visible=False)
    filename_choices = gr.update(choices=[], value=None)
    accordion_update = gr.update(label="File Content")
    code_update = gr.update(value="")

    if stego_image_pil is None:
        status = "Error: Please upload or paste an image."
        return status, file_data_state, file_buffers_state, file_browser_visibility, filename_choices, accordion_update, code_update

    if not password:
        status = "Error: Please enter a decryption password."
        return status, file_data_state, file_buffers_state, file_browser_visibility, filename_choices, accordion_update, code_update

    try:
        # Step 1: Extract and decrypt data from the image
        extracted_data = keylock_core.extract_data_from_image(stego_image_pil.convert("RGB"))
        decrypted_bytes = keylock_core.decrypt_data(extracted_data, password)
        
        # Step 2: Decode the extracted data (JSON or plain text)
        try:
            data = json.loads(decrypted_bytes.decode('utf-8'))
            md_output = data.get('repo_md', json.dumps(data, indent=2))
            status = f"Success! Data extracted for repo: {data.get('repo_name', 'N/A')}"
        except (json.JSONDecodeError, UnicodeDecodeError):
            md_output = decrypted_bytes.decode('utf-8', errors='ignore')
            status = "Warning: Decrypted data was not valid JSON, but was decoded as text."
        
        # Step 3: Process the markdown content into files
        file_data, file_buffers = markdown_to_files(md_output)
        if isinstance(file_data, str):
             status = f"Extraction successful, but file parsing failed: {file_data}"
             return status, [], {}, file_browser_visibility, filename_choices, accordion_update, code_update

        # Step 4: Prepare the UI updates
        file_data_state, file_buffers_state = file_data, file_buffers
        filenames = [f['filename'] for f in file_data]
        
        if filenames:
            file_browser_visibility = gr.update(visible=True)
            filename_choices = gr.update(choices=filenames, value=filenames[0])
            first_file = file_data[0]
            accordion_update = gr.update(label=f"File Content: {first_file['filename']}")
            code_update = gr.update(value=first_file['content'])
        else:
            status += " (No files were found in the decoded content)."

    except ValueError as e:
        status = f"Error: Failed to decrypt. Most likely an incorrect password. ({e})"
    except Exception as e:
        status = f"An unexpected error occurred: {e}"
    
    return status, file_data_state, file_buffers_state, file_browser_visibility, filename_choices, accordion_update, code_update


def update_file_preview(selected_filename, file_data):
    """Updates the file content preview when a new file is selected from the dropdown."""
    if not selected_filename or not file_data:
        return gr.update(label="File Content"), gr.update(value="")
    
    selected_file = next((f for f in file_data if f['filename'] == selected_filename), None)
    
    if selected_file:
        return gr.update(label=f"File Content: {selected_file['filename']}"), gr.update(value=selected_file['content'])
    
    return gr.update(label="File Content"), gr.update(value="Error: File not found in state.")


def download_all_zip(buffers):
    """Creates a zip file from the file buffers and returns its path for download."""
    if not buffers:
        return None
    
    with tempfile.NamedTemporaryFile(delete=False, suffix=".zip", prefix="decoded_files_") as tmp:
        with zipfile.ZipFile(tmp.name, "w", zipfile.ZIP_DEFLATED) as zf:
            for filename, content in buffers.items():
                zf.writestr(filename, content)
        return tmp.name


# --- GRADIO UI DEFINITION ---
with gr.Blocks(theme=gr.themes.Soft(), title="KeyLock Decoder") as demo:
    gr.Markdown("# KeyLock Image Decoder")
    gr.Markdown("Upload or paste your KeyLock image, enter the password, and click 'Extract Files' to view and download the contents.")
    
    # State variables to hold data between user interactions
    file_data_state = gr.State([])
    file_buffers_state = gr.State({})

    with gr.Row():
        with gr.Column(scale=1):
            # --- INPUTS ---
            extract_stego_image_upload = gr.Image(
                label="Upload or Paste KeyLock Image",
                type="pil",
                sources=["upload", "clipboard"]
            )
            extract_password_input = gr.Textbox(
                label="Decryption Password",
                type="password",
                placeholder="Enter the password used for encryption"
            )
            extract_button = gr.Button(
                value="Extract Files",
                variant="primary"
            )
            
        with gr.Column(scale=2):
            # --- OUTPUTS ---
            extract_output_status = gr.Textbox(
                label="Extraction Status",
                interactive=False,
                placeholder="Status messages will appear here..."
            )
            
            with gr.Column(visible=False) as file_browser_ui:
                gr.Markdown("--- \n ### Decoded Files")
                download_all_zip_btn = gr.Button("Download All as .zip")
                
                file_selector_dd = gr.Dropdown(
                    label="Select a file to preview",
                    interactive=True
                )
                with gr.Accordion("File Content", open=True) as file_preview_accordion:
                    file_preview_code = gr.Code(
                        language="markdown", # Will auto-detect based on file in many cases
                        interactive=False,
                        label="File Preview"
                    )
                
                # Hidden component to handle the zip download
                download_zip_output = gr.File(
                    label="Download .zip file",
                    interactive=False
                )

    # --- EVENT HANDLERS ---
    
    extract_button.click(
        fn=retrieve_and_process_files,
        inputs=[
            extract_stego_image_upload,
            extract_password_input
        ],
        outputs=[
            extract_output_status,
            file_data_state,
            file_buffers_state,
            file_browser_ui,
            file_selector_dd,
            file_preview_accordion,
            file_preview_code
        ]
    )

    file_selector_dd.change(
        fn=update_file_preview,
        inputs=[
            file_selector_dd,
            file_data_state
        ],
        outputs=[
            file_preview_accordion,
            file_preview_code
        ]
    )

    download_all_zip_btn.click(
        fn=download_all_zip,
        inputs=[file_buffers_state],
        outputs=[download_zip_output]
    )

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