Update app.py

app.py

@@ -3,178 +3,223 @@ from gradio_client import Client
 from PIL import Image
 import base64
 import io
+import json
 import logging
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
 
 # --- Configure Logging ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 # ==============================================================================
-#  CONFIGURATION: IDs of the remote Gradio Spaces
+#  CONFIGURATION: IDs AND URLs OF THE REMOTE SERVICES
 # ==============================================================================
 CREATOR_SPACE_ID = "broadfield-dev/KeyLock-Auth-Creator"
 SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server"
-CLIENT_SPACE_LINK = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-Client"
 
+# 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}"
+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
 # ==============================================================================
 
-def create_image_via_api(secret_data: str, public_key: str):
-    """
-    Calls the Creator Space API to generate an encrypted image.
-    This function uses 'yield' to provide real-time status updates to the UI.
-    """
-    if not all([secret_data, public_key]):
-        raise gr.Error("Secret Data and Public Key are both required.")
+def get_creator_endpoints():
+    """Calls the Creator space to get its list of supported endpoints."""
+    status = f"Fetching endpoint list from {CREATOR_SPACE_ID}..."
+    yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status
+    try:
+        client = Client(src=CREATOR_SPACE_ID)
+        json_string = client.predict(api_name="/get_endpoints")
+        
+        endpoints = json.loads(json_string)
+        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
+    except Exception as e:
+        logger.error(f"Failed to get endpoints from creator: {e}", exc_info=True)
+        status = f"❌ Error: Could not fetch endpoints from Creator space. Check if it's running. Details: {e}"
+        yield gr.Dropdown(choices=[], value=None, label="Error fetching services"), status
+
+def create_image_via_api(service_name: str, secret_data: str, available_endpoints: list):
+    """Calls the Creator Space API to generate an encrypted image for a selected service."""
+    if not all([service_name, secret_data]):
+        raise gr.Error("Please select a service and provide secret data.")
     
-    status = f"Initializing client for Creator: {CREATOR_SPACE_ID}..."
-    yield None, None, status  # Yield initial status update
+    status = f"Looking up public key for '{service_name}'..."
+    yield None, None, status
 
+    # Find the public key from the list we fetched earlier
+    public_key = next((e['public_key'] for e in available_endpoints if e['name'] == service_name), None)
+    if not public_key:
+        raise gr.Error(f"Could not find public key for '{service_name}' in the fetched configuration.")
+        
+    status = f"Connecting to Creator: {CREATOR_SPACE_ID}..."
+    yield None, None, status
+    
     try:
         client = Client(src=CREATOR_SPACE_ID)
-        status = f"Calling API '/create_image' on {CREATOR_SPACE_ID}..."
-        yield None, None, status
+        temp_filepath = client.predict(secret_data, public_key, api_name="/create_image")
         
-        # The 'predict' method calls the API endpoint.
-        # The result from an Image output in the remote API is a filepath to a temporary file.
-        temp_filepath = client.predict(
-            secret_data_str=secret_data,
-            public_key_pem=public_key,
-            api_name="/create_image"  # The API name defined in the Creator space
-        )
+        if not temp_filepath: raise gr.Error("Creator API did not return an image.")
         
-        if not temp_filepath:
-            raise gr.Error("Creator API did not return a valid image file path.")
-            
-        # Load the image from the temporary file to display it in our dashboard's UI
         created_image = Image.open(temp_filepath)
-        
-        status = "✅ Success! Image created by the Creator service."
+        status = f"✅ Success! Image created for '{service_name}'."
         yield created_image, temp_filepath, status
-
     except Exception as e:
         logger.error(f"Creator API call failed: {e}", exc_info=True)
-        # Yield the error message back to the UI
         yield None, None, f"❌ Error calling Creator API: {e}"
 
-
 def decrypt_image_via_api(image: Image.Image):
-    """
-    Calls the Server Space API to decrypt an image.
-    Uses 'yield' for status updates.
-    """
-    if image is None:
-        raise gr.Error("Please upload an image to decrypt.")
-
-    status = f"Initializing client for Server: {SERVER_SPACE_ID}..."
+    """Calls the Server Space API to decrypt an image."""
+    if image is None: raise gr.Error("Please upload an image to decrypt.")
+    status = f"Connecting to Server: {SERVER_SPACE_ID}..."
     yield None, status
-
     try:
         client = Client(src=SERVER_SPACE_ID)
-        
-        # Convert the user's uploaded PIL image to a base64 string for the server's API
         with io.BytesIO() as buffer:
             image.save(buffer, format="PNG")
             b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
         
-        status = f"Calling API '/keylock-auth-decoder' on {SERVER_SPACE_ID}..."
+        status = f"Calling API on {SERVER_SPACE_ID}..."
         yield None, status
-        
-        # Call the server's API endpoint. The result will be a dictionary.
-        decrypted_json = client.predict(
-            image_base64_string=b64_string,
-            api_name="/keylock-auth-decoder"  # The API name defined in the Server space
-        )
-        
+        decrypted_json = client.predict(b64_string, api_name="/keylock-auth-decoder")
         status = "✅ Success! Data decrypted by the Server."
         yield decrypted_json, status
-
     except Exception as e:
         logger.error(f"Server API call failed: {e}", exc_info=True)
         yield None, f"❌ Error calling Server API: {e}"
 
+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,
+        encryption_algorithm=serialization.NoEncryption()
+    ).decode('utf-8')
+    public_pem = private_key.public_key().public_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PublicFormat.SubjectPublicKeyInfo
+    ).decode('utf-8')
+    return private_pem, public_pem
 
 # ==============================================================================
 #  GRADIO DASHBOARD INTERFACE
 # ==============================================================================
 theme = gr.themes.Base(
-    primary_hue="blue",
-    secondary_hue="sky",
-    neutral_hue="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="#F8FAFC",
+    body_background_fill="#F1F5F9", # Light grey
+    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",
-    button_primary_text_color="white",
 )
 
 with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
+    # This State component will hold the list of dicts fetched from the Creator
+    endpoints_state = gr.State([])
+
     gr.Markdown("# 🔑 KeyLock Operations Dashboard")
-    gr.Markdown("A centralized dashboard demonstrating the entire KeyLock ecosystem by making live API calls to the deployed Creator and Server spaces.")
+    gr.Markdown("A centralized dashboard to manage and demonstrate the entire KeyLock ecosystem, powered by live API calls to dedicated services.")
 
-    with gr.Tabs():
-        with gr.TabItem("🏭 Image Creator", id=0):
-            gr.Markdown("## Create an Encrypted Image via API Call")
-            gr.Markdown(f"This interface calls the **`{CREATOR_SPACE_ID}`** Space to generate a new encrypted image.")
+    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.Column():
+                    gen_keys_button = gr.Button("Generate New 2048-bit Key Pair", icon="🔑", variant="secondary")
+                    output_private_key = gr.Textbox(lines=10, label="Generated Private Key (Keep Secret!)", interactive=False, show_copy_button=True)
+                output_public_key = gr.Textbox(lines=10, label="Generated Public Key (Share This)", interactive=False, show_copy_button=True)
+
+        with gr.TabItem("② Auth Creator", id=1):
+            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):
-                    creator_secret_input = gr.Textbox(lines=5, label="Secret Data", placeholder="API_KEY: sk-123...\nUSER: demo-user")
-                    creator_pubkey_input = gr.Textbox(lines=7, label="Public Key of the Target Server", placeholder="Paste the Server's public key here...")
-                    creator_button = gr.Button("Create Image via Creator API", variant="primary", icon="✨")
+                    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", icon="🔄", scale=0)
+                    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", icon="✨")
                 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_download_output = gr.File(label="Download Image File")
 
-        with gr.TabItem("💻 Client / Decoder", id=1):
-            gr.Markdown("## Decrypt an Image via API Call")
-            gr.Markdown(f"This interface acts as a client, calling the **`{SERVER_SPACE_ID}`** Space to decrypt an image.")
+        with gr.TabItem("③ Client / Decoder", id=2):
+            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"):
                 with gr.Column(scale=1):
-                    client_image_input = gr.Image(type="pil", label="Upload Encrypted Image", sources=["upload", "clipboard"])
+                    client_image_input = gr.Image(type="pil", label="Upload Encrypted Auth Image", sources=["upload", "clipboard"])
                     client_button = gr.Button("Decrypt Image via Server API", variant="primary", icon="🔓")
                 with gr.Column(scale=1):
                     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=2):
-            gr.Markdown("## Ecosystem Overview")
+        with gr.TabItem("ℹ️ Service Information", id=3):
+            gr.Markdown("## Ecosystem Architecture")
             gr.Markdown(
-                f"""
-                This dashboard coordinates three separate Hugging Face Spaces to demonstrate a complete workflow:
-
-                1.  **Creator Service:**
-                    -   **Space:** [{CREATOR_SPACE_ID}](https://huggingface.co/spaces/{CREATOR_SPACE_ID})
-                    -   **Role:** Provides a UI and an API (`/run/create_image`) to encrypt data into PNG images.
-
-                2.  **Server (Decoder) Service:**
-                    -   **Space:** [{SERVER_SPACE_ID}](https://huggingface.co/spaces/{SERVER_SPACE_ID})
-                    -   **Role:** Holds a secret private key and provides a secure API (`/run/keylock-auth-decoder`) to decrypt images.
-
-                3.  **This Dashboard:**
-                    -   **Role:** Acts as a master client and control panel, making live API calls to the other services to showcase the end-to-end process.
-                
-                *Note: The original `{CLIENT_SPACE_LINK.split('/')[-1]}` is now superseded by the 'Client / Decoder' tab in this dashboard.*
+                """
+                This dashboard coordinates separate Hugging Face Spaces to demonstrate a secure, decoupled workflow. Each service has a specific role.
                 """
             )
-
-    creator_button.click(
-        fn=create_image_via_api,
-        inputs=[creator_secret_input, creator_pubkey_input],
-        outputs=[creator_image_output, creator_download_output, creator_status]
-    )
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown(f"""
+                    ### 🏭 Auth Creator Service
+                    - **Space:** [{CREATOR_SPACE_ID}]({CREATOR_URL})
+                    - **Role:** Provides a UI and an API to encrypt data into PNG images for various targets. It reads a public configuration file (`endpoints.json`) to know which public keys it can use.
+                    - **Source Code:** [app.py]({CREATOR_APP_PY_URL})
+                    """)
+                with gr.Column():
+                    gr.Markdown(f"""
+                    ### 📡 Decoder Server
+                    - **Space:** [{SERVER_SPACE_ID}]({SERVER_URL})
+                    - **Role:** The trusted authority. It holds a **secret private key** and provides a secure API to decrypt images. This is the only component that can read the secret data.
+                    - **Source Code:** [app.py]({SERVER_APP_PY_URL})
+                    """)
+
+    # --- Wire up the component logic ---
+    gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
     
-    client_button.click(
-        fn=decrypt_image_via_api,
-        inputs=[client_image_input],
-        outputs=[client_json_output, client_status]
+    # Logic for fetching and populating the dropdown
+    load_event = demo.load(
+        fn=get_creator_endpoints,
+        inputs=None,
+        outputs=[creator_service_dropdown, creator_status]
+    ).then(
+        lambda data: data, # Pass the full data to the state
+        _js="(data) => data.choices",
+        outputs=[endpoints_state]
+    )
+    refresh_button.click(
+        fn=get_creator_endpoints,
+        inputs=None,
+        outputs=[creator_service_dropdown, creator_status]
+    ).then(
+        lambda data: data,
+        _js="(data) => data.choices",
+        outputs=[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])
+    client_button.click(fn=decrypt_image_via_api, inputs=[client_image_input], outputs=[client_json_output, client_status])
+
 if __name__ == "__main__":
     demo.launch()