Update app.py

app.py

@@ -12,29 +12,22 @@ from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from cryptography.hazmat.primitives import hashes
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa, padding
-from cryptography.exceptions import InvalidTag
 
 # --- Configure Logging ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 # ==============================================================================
-#  CONFIGURATION: URL of the Remote SERVER Service
+#  CONFIGURATION: URL to the public list of services
 # ==============================================================================
-SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server"
-BASE_HF_URL = "https://huggingface.co/spaces/"
-SERVER_URL = f"{BASE_HF_URL}{SERVER_SPACE_ID}"
-# The API endpoint is constructed from the server's direct URL, not the hub URL.
-SERVER_DIRECT_URL_BASE = f"https://{SERVER_SPACE_ID.replace('/', '-')}.hf.space"
-SERVER_API_ENDPOINT = f"{SERVER_DIRECT_URL_BASE}/run/keylock-auth-decoder"
+# This can be any publicly accessible raw JSON file.
+ENDPOINTS_JSON_URL = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-Creator/raw/main/endpoints.json"
 
 # ==============================================================================
 #  LOCAL LOGIC (Key and Image Generation)
 # ==============================================================================
-
 def generate_rsa_keys():
     """Generates a new 2048-bit RSA key pair LOCALLY."""
-    logger.info("Generating new RSA key pair locally.")
     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,
@@ -47,7 +40,6 @@ def generate_rsa_keys():
 
 def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.Image:
     """Creates the encrypted image LOCALLY."""
-    logger.info("Starting local image creation process...")
     if not secret_data_str.strip(): raise ValueError("Secret data cannot be empty.")
     if not public_key_pem.strip(): raise ValueError("Public Key cannot be empty.")
     
@@ -78,19 +70,50 @@ def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.I
     binary_payload = ''.join(format(byte, '08b') for byte in struct.pack('>I', len(encrypted_payload)) + encrypted_payload)
     if len(binary_payload) > pixel_data.size: raise ValueError("Data too large for image capacity.")
     for i in range(len(binary_payload)): pixel_data[i] = (pixel_data[i] & 0xFE) | int(binary_payload[i])
-    
     stego_pixels = pixel_data.reshape((600, 800, 3))
     return Image.fromarray(stego_pixels, 'RGB')
 
 # ==============================================================================
-#  REMOTE API CALL LOGIC
+#  UI HELPER & REMOTE API CALL LOGIC
 # ==============================================================================
 
-def decrypt_image_via_api(image: Image.Image):
-    """Makes a LIVE API call to the deployed server to decrypt an image."""
-    if image is None: raise gr.Error("Please provide an image to send.")
-        
-    status = f"Connecting to server: {SERVER_SPACE_ID}..."
+def get_server_list():
+    """Fetches the list of servers from the public JSON file."""
+    status = f"Fetching server list from remote config..."
+    yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, []
+    try:
+        response = requests.get(ENDPOINTS_JSON_URL, timeout=10)
+        response.raise_for_status()
+        endpoints = response.json()
+        endpoint_names = [e['name'] for e in endpoints]
+        status = f"✅ Success! Found {len(endpoint_names)} servers."
+        yield gr.Dropdown(choices=endpoint_names, value=endpoint_names[0] if endpoint_names else None, label="Target Server"), status, endpoints
+    except Exception as e:
+        status = f"❌ Error fetching configuration: {e}"
+        yield gr.Dropdown(choices=[], value=None, label="Error fetching servers"), status, []
+
+def create_keylock_wrapper(service_name: str, secret_data: str, available_endpoints: list):
+    """UI wrapper for creating an image for a selected service."""
+    if not service_name: raise gr.Error("Please select a target server.")
+    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}'.")
+    try:
+        created_image = create_encrypted_image(secret_data, public_key)
+        return created_image, f"✅ Success! Image created for '{service_name}'."
+    except Exception as e:
+        return None, f"❌ Error: {e}"
+
+def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoints: list):
+    """UI wrapper for sending an image to a selected server API."""
+    if not service_name: raise gr.Error("Please select a target server.")
+    if image is None: raise gr.Error("Please upload an image to send.")
+    
+    server = next((e for e in available_endpoints if e['name'] == service_name), None)
+    if not server: raise gr.Error(f"Could not find server configuration for '{service_name}'.")
+
+    server_space_id = server.get('link', '').split('/spaces/')[-1]
+    api_endpoint = f"https://{server_space_id.replace('/', '-')}.hf.space/run/keylock-auth-decoder"
+    status = f"Connecting to server: {server_space_id}..."
     yield None, status
     
     try:
@@ -100,31 +123,19 @@ def decrypt_image_via_api(image: Image.Image):
         
         payload = {"data": [b64_string]}
         headers = {"Content-Type": "application/json"}
-        
-        status = f"Sending image to API endpoint:\n{SERVER_API_ENDPOINT}"
-        yield None, status
-        
-        response = requests.post(SERVER_API_ENDPOINT, headers=headers, json=payload, timeout=45)
+        response = requests.post(api_endpoint, headers=headers, json=payload, timeout=45)
         response_json = response.json()
         
         if response.status_code == 200:
             if "data" in response_json:
-                decrypted_data = response_json["data"][0]
-                status = "✅ Success! Data decrypted by the remote server."
-                return decrypted_data, status
-            elif "error" in response_json:
-                raise gr.Error(f"API returned an error: {response_json['error']}")
+                return response_json["data"][0], "✅ Success! Data decrypted by remote server."
+            else:
+                raise gr.Error(f"API returned an unexpected success format: {response_json}")
         else:
-            error_detail = response_json.get("error", "Unknown error.")
-            raise gr.Error(f"API Error (Status {response.status_code}): {error_detail}")
-
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Network error calling API: {e}")
-        raise gr.Error(f"Could not connect to the API. Check the server space is running and the URL is correct. Error: {e}")
+            raise gr.Error(f"API Error (Status {response.status_code}): {response_json.get('error', 'Unknown error')}")
+            
     except Exception as e:
-        logger.error(f"An unexpected error occurred: {e}", exc_info=True)
-        raise gr.Error(f"An unexpected error occurred: {e}")
-
+        return None, f"❌ Error calling server API: {e}"
 
 # ==============================================================================
 #  GRADIO DASHBOARD INTERFACE
@@ -139,83 +150,67 @@ theme = gr.themes.Base(
 )
 
 with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
+    endpoints_state = gr.State([])
+
     gr.Markdown("# 🔑 KeyLock Operations Dashboard")
     gr.Markdown("A self-contained dashboard to demonstrate the KeyLock ecosystem. Key/Image creation is performed locally, while decryption is handled by a **live, remote API call** to a secure server.")
 
     with gr.Tabs() as tabs:
-        with gr.TabItem("① Generate Keys", id=0):
-            gr.Markdown("## Step 1: Create a Secure Key Pair (Local)")
-            gr.Markdown(
-                """
-                This tool generates a new RSA key pair within your browser session. In a real-world scenario, the **Private Key** would be immediately stored as a secure secret on a server (like the `KEYLOCK_PRIV_KEY` secret on our demo server), and would never be shown in a UI like this. The **Public Key** would be distributed to clients or other services that need to encrypt data for that server.
-                
-                **Action:** Click the button below, then copy both keys for the next steps.
-                """
-            )
+        with gr.TabItem("① Create KeyLock", id=0):
+            gr.Markdown("## Step 1: Create an Encrypted Authentication Image (Local)")
+            gr.Markdown(f"This tool acts as the **Auth Creator**. It fetches a list of available servers from a public [configuration file]({ENDPOINTS_JSON_URL}), then uses the selected server's public key to encrypt your data into a PNG. **This entire creation process happens locally.**")
             with gr.Row(variant="panel"):
-                with gr.Column(scale=1):
-                    gr.Markdown("### Your New Keys")
-                    gen_keys_button = gr.Button("Generate New 2048-bit Key Pair", icon="🔑", variant="secondary")
                 with gr.Column(scale=2):
-                     with gr.Row():
-                        output_public_key = gr.Textbox(lines=11, label="Generated Public Key (For Creator)", interactive=False, show_copy_button=True)
-                        output_private_key = gr.Textbox(lines=11, label="Generated Private Key (For Decoder)", interactive=False, show_copy_button=True)
-
-        with gr.TabItem("② Create KeyLock", id=1):
-            gr.Markdown("## Step 2: Create an Encrypted Auth Image (Local)")
-            gr.Markdown(
-                """
-                This tool acts as the **Auth Creator**. It takes your secret data and uses the **Public Key** you generated in Step 1 to encrypt it into a new PNG image. This entire process happens locally in this application. This simulates a user or an automated client preparing credentials to send to the secure server.
-
-                **Action:** Paste the **Public Key** from Step 1, enter some secrets, and click create.
-                """
-            )
-            with gr.Row(variant="panel"):
-                with gr.Column(scale=1):
-                    gr.Markdown("### Configuration")
-                    creator_pubkey_input = gr.Textbox(lines=8, label="Paste the Public Key Here", placeholder="Copy the public key generated in Step 1...")
-                    creator_secret_input = gr.Textbox(lines=5, label="Secret Data to Encrypt", placeholder="SESSION_ID: abc-123\nUSER: [email protected]")
+                    with gr.Row():
+                        creator_service_dropdown = gr.Dropdown(label="Target Server", interactive=True, info="Select the API server to encrypt data for.")
+                        refresh_button = gr.Button("🔄", scale=0, size="sm", tooltip="Refresh Server List from Config File")
+                    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", variant="primary")
                 with gr.Column(scale=1):
-                    gr.Markdown("### Output")
                     creator_status = gr.Textbox(label="Status", interactive=False, lines=2)
-                    creator_image_output = gr.Image(label="Generated Encrypted Image", type="pil", show_download_button=True, format="png", show_share_button=False)
-
-        with gr.TabItem("③ Send KeyLock", id=2):
-            gr.Markdown("## Step 3: Decrypt via Live API Call")
-            gr.Markdown(
-                f"""
-                This is the core demonstration. This tool acts as a **Client** sending the encrypted image to our live, remote **Server** at [{SERVER_SPACE_ID}]({SERVER_URL}). 
-                
-                For this demo to work, the **Private Key** you generated in Step 1 must be **the same one** set as the `KEYLOCK_PRIV_KEY` secret in the `{SERVER_SPACE_ID}` Space settings. The client sends the image, and the server uses its own secret key to decrypt it.
-
-                **Action:** Upload the image from Step 2. The dashboard will make a live API call to `{SERVER_API_ENDPOINT}`.
-                """
-            )
+                    creator_image_output = gr.Image(label="Generated Encrypted Image", type="pil", show_download_button=True, format="png")
+
+        with gr.TabItem("② Send KeyLock", id=1):
+            gr.Markdown("## Step 2: Decrypt via Live API Call")
+            gr.Markdown("This tool acts as the **Client**. It sends the encrypted image you created in Step 1 to the live, remote **Decoder Server** you select. The server uses its securely stored private key to decrypt the data and sends the result back.")
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1):
-                    gr.Markdown("### Input")
+                    gr.Markdown("### Select Server to Call")
+                    send_service_dropdown = gr.Dropdown(label="Target Server", interactive=True, info="Select the API server to send the image to.")
+                    gr.Markdown("### Upload Image")
                     client_image_input = gr.Image(type="pil", label="Upload or Drag Encrypted Image Here", sources=["upload", "clipboard"])
-                    client_button = gr.Button("🔓 Decrypt Image via Remote Server", variant="primary")
+                    client_button = gr.Button("🔓 Decrypt via Remote Server", variant="primary")
                 with gr.Column(scale=1):
-                    gr.Markdown("### Decrypted Data")
+                    gr.Markdown("### Response from Server")
                     client_status = gr.Textbox(label="Status", interactive=False, lines=2)
-                    client_json_output = gr.JSON(label="Result from Server")
+                    client_json_output = gr.JSON(label="Decrypted Data")
+        
+        with gr.TabItem("ℹ️ Key Generation & Info", id=2):
+            gr.Markdown("## Ecosystem Information")
+            gr.Markdown("This dashboard coordinates with live Hugging Face Spaces to demonstrate a secure, decoupled workflow.")
+            with gr.Accordion("🔑 RSA Key Pair Generator", open=False):
+                gr.Markdown("Create a new public/private key pair. In a real scenario, you would add the **Public Key** to the public `endpoints.json` configuration file, and set the **Private Key** as a secret variable in the corresponding server space.")
+                with gr.Row():
+                    with gr.Column():
+                        output_public_key = gr.Textbox(lines=10, label="Generated Public Key", interactive=False, show_copy_button=True)
+                    with gr.Column():
+                        output_private_key = gr.Textbox(lines=10, label="Generated Private Key", 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])
     
-    creator_button.click(
-        fn=create_encrypted_image,
-        inputs=[creator_pubkey_input, creator_secret_input],
-        outputs=[creator_image_output, creator_status]
-    )
-    
-    client_button.click(
-        fn=decrypt_image_via_api,
-        inputs=[client_image_input],
-        outputs=[client_json_output, client_status]
-    )
+    def refresh_and_update_all():
+        *_, last_yield = get_server_list()
+        dropdown_update, status_update, state_update = last_yield
+        # Update both dropdowns simultaneously
+        return dropdown_update, dropdown_update, status_update, state_update
+
+    refresh_button.click(fn=refresh_and_update_all, outputs=[creator_service_dropdown, send_service_dropdown, creator_status, endpoints_state])
+    demo.load(fn=refresh_and_update_all, outputs=[creator_service_dropdown, send_service_dropdown, creator_status, endpoints_state])
+
+    creator_button.click(fn=create_keylock_wrapper, inputs=[creator_service_dropdown, creator_secret_input, endpoints_state], outputs=[creator_image_output, creator_status])
+    client_button.click(fn=send_keylock_wrapper, inputs=[send_service_dropdown, client_image_input, endpoints_state], outputs=[client_json_output, client_status])
 
 if __name__ == "__main__":
     demo.launch()