Update app.py

app.py

@@ -1,5 +1,4 @@
 import gradio as gr
-from gradio_client import Client
 from PIL import Image, ImageDraw, ImageFont
 import base64
 import io
@@ -7,7 +6,6 @@ import json
 import logging
 import os
 import requests
-import tempfile
 import struct
 import numpy as np
 from cryptography.hazmat.primitives import serialization
@@ -20,20 +18,18 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(level
 logger = logging.getLogger(__name__)
 
 # ==============================================================================
-#  CONFIGURATION: IDs AND URLs OF THE REMOTE SERVICES
+#  CONFIGURATION
 # ==============================================================================
-CREATOR_SPACE_ID = "broadfield-dev/KeyLock-Auth-Creator"
-SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server"
 CREATOR_ENDPOINTS_JSON_URL = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-Creator/raw/main/endpoints.json"
-
 BASE_HF_URL = "https://huggingface.co/spaces/"
+# Define these here for cleaner linking in the UI
+CREATOR_SPACE_ID = "broadfield-dev/KeyLock-Auth-Creator"
+SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server"
 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"
 
 # ==============================================================================
-#  LOCAL LOGIC (Key and Image Generation) - from your working version
+#  LOCAL LOGIC (Key and Image Generation)
 # ==============================================================================
 def generate_rsa_keys():
     private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
@@ -73,31 +69,35 @@ def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.I
 #  UI HELPER & REMOTE API CALL LOGIC
 # ==============================================================================
 
-# --- THIS IS THE MISSING FUNCTION THAT HAS BEEN ADDED BACK ---
 def get_server_list():
-    """Fetches the list of servers from the public JSON file."""
+    """Fetches and validates the list of servers from the public JSON file."""
     status = f"Fetching server list from remote config..."
-    # Yield an initial state to show the UI is working
     yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, []
     try:
-        response = requests.get(CREATOR_ENDPOINTS_JSON_URL, timeout=10)
+        response = requests.get(ENDPOINTS_JSON_URL, timeout=10)
         response.raise_for_status()
-        endpoints = response.json()
-        # Validate the structure to prevent downstream errors
-        for e in endpoints:
-            if not all(k in e for k in ["name", "api_endpoint", "public_key"]):
-                raise ValueError(f"Configuration entry is missing required keys: {e}")
-        endpoint_names = [e['name'] for e in endpoints]
-        status = f"✅ Success! Found {len(endpoint_names)} servers."
-        # Yield the final, successful state
-        yield gr.Dropdown(choices=endpoint_names, value=endpoint_names[0] if endpoint_names else None, label="Target Server"), status, endpoints
+        
+        all_entries = response.json()
+        valid_endpoints = []
+        required_keys = ["name", "api_endpoint", "public_key"]
+        
+        # --- ROBUST VALIDATION LOGIC ---
+        for entry in all_entries:
+            if all(key in entry for key in required_keys):
+                valid_endpoints.append(entry)
+            else:
+                logger.warning(f"Skipping invalid entry in configuration file: {entry}")
+
+        if not valid_endpoints:
+            raise ValueError("No valid server configurations found in the remote file.")
+
+        endpoint_names = [e['name'] for e in valid_endpoints]
+        status = f"✅ Success! Found {len(endpoint_names)} valid servers."
+        yield gr.Dropdown(choices=endpoint_names, value=endpoint_names[0] if endpoint_names else None, label="Target Server"), status, valid_endpoints
     except Exception as e:
-        logger.error(f"Failed to get server list: {e}", exc_info=True)
-        status = f"❌ Error fetching configuration: {e}"
-        # Yield a final error state
+        status = f"❌ Error fetching or parsing 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.")
@@ -105,7 +105,6 @@ def create_keylock_wrapper(service_name: str, secret_data: str, available_endpoi
     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)
-        # Use a single return for robustness
         return created_image, f"✅ Success! Image created for '{service_name}'."
     except Exception as e:
         return None, f"❌ Error: {e}"
@@ -116,11 +115,9 @@ def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoi
     if image is None: raise gr.Error("Please upload an image to send.")
     
     api_endpoint = next((e.get('api_endpoint') for e in available_endpoints if e['name'] == service_name), None)
-    if not api_endpoint:
-        raise gr.Error(f"Configuration Error: Could not find 'api_endpoint' for '{service_name}'.")
+    if not api_endpoint: raise gr.Error(f"Configuration Error: Could not find 'api_endpoint' for '{service_name}'.")
 
     status = f"Connecting to endpoint: {api_endpoint}"
-    # Use yield to provide an immediate status update
     yield None, status
     
     try:
@@ -211,11 +208,9 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
     gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
     
     def refresh_and_update_all():
-        # The generator must be iterated to get its values
+        # The generator must be iterated to get its final values
         for dropdown_update, status_update, state_update in get_server_list():
-            # This loop will run twice: once for the "loading" state, once for the final state.
-            pass # We only care about the final state
-        # Update both dropdowns and the state
+            pass
         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])