Update keylock_component.py

keylock_component.py

@@ -15,7 +15,7 @@ import numpy as np
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 
-# --- Backend Logic Class (Updated to use tempfile) ---
+# --- Backend Logic Class ---
 class AppServerLogic:
     """Contains all backend logic and mock data needed by the application."""
     def __init__(self):
@@ -62,10 +62,7 @@ class AppServerLogic:
             return {"status": "Error", "message": f"Decryption Failed: {e}"}
             
     def generate_encrypted_image(self, payload_dict):
-        """
-        Creates a KeyLock image and returns both a Pillow object for preview
-        and a filepath to a saved, uncorrupted PNG for download.
-        """
+        """Creates a KeyLock image and returns its filepath for both preview and download."""
         try:
             response = requests.get("https://images.unsplash.com/photo-1506318137071-a8e063b4bec0?q=80&w=1200", timeout=10)
             img = ImageOps.fit(Image.open(io.BytesIO(response.content)).convert("RGB"), (600, 600))
@@ -84,11 +81,10 @@ class AppServerLogic:
         pixel_data[:len(binary_payload)] = (pixel_data[:len(binary_payload)] & 0xFE) | np.array(list(binary_payload), dtype=np.uint8)
         final_image = Image.fromarray(pixel_data.reshape(img.size[1], img.size[0], 3), 'RGB')
         
-        # Save the pristine PNG to a temporary file on the server.
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
             final_image.save(f.name, "PNG")
-            # Return both the Pillow image for preview and the string filepath for download.
-            return final_image, f.name
+            # Return the filepath for both the preview image and the download file.
+            return f.name, f.name
 
     @staticmethod
     def generate_pem_keys():
@@ -97,7 +93,7 @@ class AppServerLogic:
         pub = pk.public_key().public_bytes(encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo).decode()
         return priv, pub
 
-# --- UI Component Class (Enhanced with download functionality) ---
+# --- UI Component Class ---
 class KeylockDecoderComponent(gr.components.Component):
     """A Gradio Component that decodes a KeyLock image and provides advanced tools."""
     EVENTS = ["change"]
@@ -132,15 +128,16 @@ class KeylockDecoderComponent(gr.components.Component):
                 gr.Markdown("Create a test image using the site's public key.")
                 payload_input = gr.JSON(label="Payload to Encrypt", value={"API_KEY": "sk-12345-abcde", "USER": "demo-user"})
                 generate_img_button = gr.Button("Generate Image", variant="secondary")
-                generated_image_preview = gr.Image(label="Generated Image Preview", interactive=False)
+                # This component now expects a filepath to display the uncorrupted PNG.
+                generated_image_preview = gr.Image(label="Generated Image Preview", type="filepath", interactive=False)
                 generated_file_download = gr.File(label="Download Uncorrupted PNG", interactive=False)
             
             with gr.Accordion("Create New Standalone Key Pair", open=False):
                 gr.Markdown("Generate a new, random key pair for testing or other uses.")
                 generate_keys_button = gr.Button("Generate Keys", variant="secondary")
                 with gr.Row():
-                    output_private_key = gr.Code(label="Generated Private Key", language="python")
-                    output_public_key = gr.Code(label="Generated Public Key", language="python")
+                    output_private_key = gr.Code(label="Generated Private Key", language="pem")
+                    output_public_key = gr.Code(label="Generated Public Key", language="pem")
 
         def on_image_upload(image):
             if image is None: return None, "Upload a KeyLock image to auto-fill credentials."