Create app.py

app.py

@@ -0,0 +1,197 @@
+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)