Update app.py

app.py

@@ -5,7 +5,8 @@ import base64
 import io
 import json
 import logging
-import requests  # Use requests to fetch the remote JSON file
+import requests
+import tempfile
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
 
@@ -18,12 +19,8 @@ logger = logging.getLogger(__name__)
 # ==============================================================================
 CREATOR_SPACE_ID = "broadfield-dev/KeyLock-Auth-Creator"
 SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server"
-
-# URL to the raw JSON file containing the list of public keys and services.
-# This makes the Creator's configuration publicly readable.
 CREATOR_ENDPOINTS_JSON_URL = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-Creator/raw/main/endpoints.json"
 
-# Construct URLs for linking in documentation
 BASE_HF_URL = "https://huggingface.co/spaces/"
 CREATOR_URL = f"{BASE_HF_URL}{CREATOR_SPACE_ID}"
 SERVER_URL = f"{BASE_HF_URL}{SERVER_SPACE_ID}"
@@ -31,7 +28,7 @@ CREATOR_APP_PY_URL = f"{CREATOR_URL}/blob/main/app.py"
 SERVER_APP_PY_URL = f"{SERVER_URL}/blob/main/app.py"
 
 # ==============================================================================
-#  API CALL WRAPPER FUNCTIONS
+#  API CALL WRAPPER FUNCTIONS (Your working versions)
 # ==============================================================================
 
 def get_creator_endpoints():
@@ -40,13 +37,10 @@ def get_creator_endpoints():
     yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, [] # Initial state
     try:
         response = requests.get(CREATOR_ENDPOINTS_JSON_URL, timeout=10)
-        response.raise_for_status()  # Raise an exception for bad status codes
-        
+        response.raise_for_status()
         endpoints = response.json()
         endpoint_names = [e['name'] for e in endpoints]
-        
         status = f"✅ Success! Found {len(endpoint_names)} endpoints."
-        # Return the full list to the state, and the updated dropdown
         yield gr.Dropdown(choices=endpoint_names, value=endpoint_names[0] if endpoint_names else None, label="Target Service"), status, endpoints
     except Exception as e:
         logger.error(f"Failed to get endpoints from creator's JSON file: {e}", exc_info=True)
@@ -74,9 +68,18 @@ def create_image_via_api(service_name: str, secret_data: str, available_endpoint
         
         if not temp_filepath: raise gr.Error("Creator API did not return an image.")
         
+        # --- PNG FIX ---
+        # The API returns a path to a temp file which could be any format.
+        # We open it, ensure it's saved as PNG, and return that path.
         created_image = Image.open(temp_filepath)
+        
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as png_file:
+            created_image.save(png_file.name, "PNG")
+            png_filepath = png_file.name
+            
         status = f"✅ Success! Image created for '{service_name}'."
-        yield created_image, temp_filepath, status
+        yield created_image, png_filepath, status
+
     except Exception as e:
         logger.error(f"Creator API call failed: {e}", exc_info=True)
         yield None, None, f"❌ Error calling Creator API: {e}"
@@ -105,13 +108,11 @@ def generate_rsa_keys():
     """Generates a new RSA key pair."""
     private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
     private_pem = private_key.private_bytes(
-        encoding=serialization.Encoding.PEM,
-        format=serialization.PrivateFormat.PKCS8,
+        encoding=serialization.Encoding.PEM, format=serialization.PrivateFormat.PKCS8,
         encryption_algorithm=serialization.NoEncryption()
     ).decode('utf-8')
     public_pem = private_key.public_key().public_bytes(
-        encoding=serialization.Encoding.PEM,
-        format=serialization.PublicFormat.SubjectPublicKeyInfo
+        encoding=serialization.Encoding.PEM, format=serialization.PublicFormat.SubjectPublicKeyInfo
     ).decode('utf-8')
     return private_pem, public_pem
 
@@ -119,18 +120,12 @@ def generate_rsa_keys():
 #  GRADIO DASHBOARD INTERFACE
 # ==============================================================================
 theme = gr.themes.Base(
-    primary_hue=gr.themes.colors.blue,
-    secondary_hue=gr.themes.colors.sky,
-    neutral_hue=gr.themes.colors.slate,
+    primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky, neutral_hue=gr.themes.colors.slate,
     font=(gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"),
 ).set(
-    body_background_fill="#F1F5F9",
-    panel_background_fill="white",
-    block_background_fill="white",
-    block_border_width="1px",
-    block_shadow="*shadow_drop_lg",
-    button_primary_background_fill="*primary_600",
-    button_primary_background_fill_hover="*primary_700",
+    body_background_fill="#F1F5F9", panel_background_fill="white", block_background_fill="white",
+    block_border_width="1px", block_shadow="*shadow_drop_lg",
+    button_primary_background_fill="*primary_600", button_primary_background_fill_hover="*primary_700",
 )
 
 with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
@@ -140,31 +135,24 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
     gr.Markdown("A centralized dashboard to manage and demonstrate the entire KeyLock ecosystem, powered by live API calls to dedicated services.")
 
     with gr.Tabs() as tabs:
-        with gr.TabItem("① Generate Keys", id=0):
-            gr.Markdown("## RSA Key Pair Generator")
-            gr.Markdown("Create a new public/private key pair. The public key can be added to a service's configuration to allow it to be a target for the Auth Creator.")
-            with gr.Row(variant="panel"):
-                with gr.Group():
-                    gen_keys_button = gr.Button("🔑 Generate New 2048-bit Key Pair", variant="secondary")
-                    with gr.Row():
-                        output_public_key = gr.Textbox(lines=10, label="Generated Public Key (Share This)", interactive=False, show_copy_button=True)
-                        output_private_key = gr.Textbox(lines=10, label="Generated Private Key (Keep Secret!)", interactive=False, show_copy_button=True)
-
-        with gr.TabItem("② Auth Creator", id=1):
+        with gr.TabItem("① Auth Creator", id=0):
             gr.Markdown("## Create an Encrypted Authentication Image")
             gr.Markdown(f"This tool calls the **[{CREATOR_SPACE_ID}]({CREATOR_URL})** service to encrypt data for a chosen target. The list of targets is fetched live from the Creator's configuration.")
             with gr.Row(variant="panel"):
                 with gr.Column(scale=2):
                     with gr.Row():
                         creator_service_dropdown = gr.Dropdown(label="Target Service", interactive=True, info="Select the API server you want to encrypt data for.")
-                        refresh_button = gr.Button("🔄 Refresh List", scale=0, size="sm")
+                        refresh_button = gr.Button("🔄", scale=0, size="sm", tooltip="Refresh Target Service List")
                     creator_secret_input = gr.Textbox(lines=8, label="Secret Data to Encrypt", placeholder="API_KEY: sk-123...\nUSER: demo-user")
                     creator_button = gr.Button("✨ Create Auth Image via API", variant="primary")
                 with gr.Column(scale=1):
                     creator_status = gr.Textbox(label="Status", interactive=False, lines=2)
-                    creator_image_output = gr.Image(label="Image from Creator Service", type="pil", show_download_button=True)
+                    creator_image_output = gr.Image(label="Image from Creator Service", type="pil", show_download_button=False) # Download handled by gr.File
+                    # --- PNG FIX ---
+                    # Use a dedicated, visible File component for reliable PNG downloads.
+                    creator_download_output = gr.File(label="Download Image as PNG", interactive=False)
 
-        with gr.TabItem("③ Client / Decoder", id=2):
+        with gr.TabItem("② Client / Decoder", id=1):
             gr.Markdown("## Decrypt an Authentication Image")
             gr.Markdown(f"This tool acts as a client, calling the **[{SERVER_SPACE_ID}]({SERVER_URL})** service to decrypt an image using its securely stored private key.")
             with gr.Row(variant="panel"):
@@ -175,7 +163,7 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
                     client_status = gr.Textbox(label="Status", interactive=False, lines=2)
                     client_json_output = gr.JSON(label="Decrypted Data from Server")
         
-        with gr.TabItem("ℹ️ Service Information", id=3):
+        with gr.TabItem("ℹ️ Service Information", id=2):
             gr.Markdown("## Ecosystem Architecture")
             gr.Markdown("This dashboard coordinates separate Hugging Face Spaces to demonstrate a secure, decoupled workflow. Each service has a specific role.")
             with gr.Row():
@@ -183,21 +171,32 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
                     gr.Markdown(f"### 🏭 Auth Creator Service\n- **Space:** [{CREATOR_SPACE_ID}]({CREATOR_URL})\n- **Role:** Provides an API to encrypt data for various targets defined in its `endpoints.json` file.\n- **Source Code:** [app.py]({CREATOR_APP_PY_URL})")
                 with gr.Column():
                     gr.Markdown(f"### 📡 Decoder Server\n- **Space:** [{SERVER_SPACE_ID}]({SERVER_URL})\n- **Role:** The trusted authority. It holds a secret private key and provides a secure API to decrypt images.\n- **Source Code:** [app.py]({SERVER_APP_PY_URL})")
+            
+            # --- Generate Keys moved to a closed Accordion ---
+            with gr.Accordion("🔑 RSA Key Pair Generator", open=False):
+                gr.Markdown("Create a new public/private key pair. The public key can be added to a service's configuration to allow it to be a target for the Auth Creator.")
+                with gr.Row():
+                    with gr.Column():
+                        output_public_key = gr.Textbox(lines=10, label="Generated Public Key (Share This)", interactive=False, show_copy_button=True)
+                    with gr.Column():
+                        output_private_key = gr.Textbox(lines=10, label="Generated Private Key (Keep Secret!)", interactive=False, show_copy_button=True)
+                gen_keys_button = gr.Button("Generate New 2048-bit Key Pair", variant="secondary")
 
     # --- Wire up the component logic ---
     gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
     
-    # Event handler for loading the page or refreshing the endpoint list
     def refresh_endpoints():
-        # This is a generator function, so we need to iterate to get the last value.
         *_, last_yield = get_creator_endpoints()
         return last_yield
 
     refresh_button.click(fn=refresh_endpoints, outputs=[creator_service_dropdown, creator_status, endpoints_state])
     demo.load(fn=refresh_endpoints, outputs=[creator_service_dropdown, creator_status, endpoints_state])
 
-    # Logic for the Creator and Client tabs
-    creator_button.click(fn=create_image_via_api, inputs=[creator_service_dropdown, creator_secret_input, endpoints_state], outputs=[creator_image_output, gr.File(visible=False), creator_status])
+    creator_button.click(
+        fn=create_image_via_api,
+        inputs=[creator_service_dropdown, creator_secret_input, endpoints_state],
+        outputs=[creator_image_output, creator_download_output, creator_status] # Now outputs to the gr.File component
+    )
     client_button.click(fn=decrypt_image_via_api, inputs=[client_image_input], outputs=[client_json_output, client_status])
 
 if __name__ == "__main__":