Update app.py

app.py

@@ -24,11 +24,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.
 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}"
@@ -36,79 +33,70 @@ 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)
+#  LOCAL LOGIC (Key and Image Generation) - from your working version
 # ==============================================================================
 def generate_rsa_keys():
-    """Generates a new 2048-bit 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,
-        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')
+    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
 
 def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.Image:
-    """Creates the encrypted image LOCALLY."""
     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.")
-    
-    # --- THIS IS THE FIX ---
-    # Replaced the failing dictionary comprehension with a standard, readable for loop.
     data_dict = {}
     for line in secret_data_str.splitlines():
         line = line.strip()
-        if not line or line.startswith('#'):
-            continue
-        
+        if not line or line.startswith('#'): continue
         parts = line.split(':', 1) if ':' in line else line.split('=', 1)
-        if len(parts) == 2:
-            key = parts[0].strip()
-            value = parts[1].strip().strip("'\"")
-            data_dict[key] = value
-            
+        if len(parts) == 2: data_dict[parts[0].strip()] = parts[1].strip().strip("'\"")
     if not data_dict: raise ValueError("No valid key-value pairs found.")
-    
     json_bytes = json.dumps(data_dict).encode('utf-8')
     public_key = serialization.load_pem_public_key(public_key_pem.encode('utf-8'))
     aes_key, nonce = os.urandom(32), os.urandom(12)
     ciphertext = AESGCM(aes_key).encrypt(nonce, json_bytes, None)
     rsa_encrypted_key = public_key.encrypt(aes_key, padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None))
     encrypted_payload = struct.pack('>I', len(rsa_encrypted_key)) + rsa_encrypted_key + nonce + ciphertext
-    
     img = Image.new('RGB', (800, 600), color=(45, 52, 54))
     draw = ImageDraw.Draw(img)
     try: font = ImageFont.truetype("DejaVuSans.ttf", 40)
     except IOError: font = ImageFont.load_default(size=30)
     draw.text((400, 300), "KeyLock Secure Data", fill=(223, 230, 233), font=font, anchor="ms")
-    
     pixel_data = np.array(img.convert("RGB")).ravel()
     binary_payload = ''.join(format(b, '08b') for b in struct.pack('>I', len(encrypted_payload)) + encrypted_payload)
     if len(binary_payload) > pixel_data.size: raise ValueError("Data too large for image.")
     pixel_data[:len(binary_payload)] = (pixel_data[:len(binary_payload)] & 0xFE) | np.array(list(binary_payload), dtype=np.uint8)
-    
     stego_pixels = pixel_data.reshape((600, 800, 3))
     return Image.fromarray(stego_pixels, 'RGB')
 
 # ==============================================================================
-#  UI HELPER & REMOTE API CALL LOGIC (Your working versions)
+#  UI HELPER & REMOTE API CALL LOGIC
 # ==============================================================================
-def get_creator_endpoints():
-    """Fetches the list of supported endpoints by making an HTTP request to the Creator's JSON file."""
-    status = f"Fetching endpoint list from {CREATOR_ENDPOINTS_JSON_URL}..."
-    yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, [] # Initial state
+
+# --- THIS IS THE MISSING FUNCTION THAT HAS BEEN ADDED BACK ---
+def get_server_list():
+    """Fetches 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.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)} endpoints."
-        yield gr.Dropdown(choices=endpoint_names, value=endpoint_names[0] if endpoint_names else None, label="Target Service"), status, 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
     except Exception as e:
-        status = f"❌ Error: Could not fetch configuration. Check the URL and if the 'endpoints.json' file is public. Details: {e}"
-        yield gr.Dropdown(choices=[], value=None, label="Error fetching services"), status, []
+        logger.error(f"Failed to get server list: {e}", exc_info=True)
+        status = f"❌ Error fetching configuration: {e}"
+        # Yield a final error state
+        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."""
@@ -117,6 +105,7 @@ 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}"
@@ -131,6 +120,7 @@ def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoi
         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:
@@ -145,17 +135,17 @@ def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoi
         
         if response.status_code == 200:
             if "data" in response_json:
-                return response_json["data"][0], "✅ Success! Data decrypted by remote server."
+                yield response_json["data"][0], "✅ Success! Data decrypted by remote server."
             else:
                 raise gr.Error(f"API returned an unexpected success format: {response_json}")
         else:
             raise gr.Error(f"API Error (Status {response.status_code}): {response_json.get('error', 'Unknown error')}")
             
     except Exception as e:
-        return None, f"❌ Error calling server API: {e}"
+        yield None, f"❌ Error calling server API: {e}"
 
 # ==============================================================================
-#  GRADIO DASHBOARD INTERFACE (Your working layout)
+#  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,
@@ -174,13 +164,14 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
 
     with gr.Tabs() as tabs:
         with gr.TabItem("① Create KeyLock", id=0):
+            # ... UI layout is the same as your working version ...
             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]({CREATOR_ENDPOINTS_JSON_URL}), then uses the selected server's public key to encrypt your data into a PNG. **This entire creation process happens locally.**")
+            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=2):
                     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")
+                        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):
@@ -188,6 +179,7 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
                     creator_image_output = gr.Image(label="Generated Encrypted Image", type="pil", show_download_button=True, format="png")
 
         with gr.TabItem("② Send KeyLock", id=1):
+            # ... UI layout is the same as your working version ...
             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 from the same configuration list. The server uses its securely stored private key to decrypt the data and sends the result back.")
             with gr.Row(variant="panel"):
@@ -203,8 +195,9 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
                     client_json_output = gr.JSON(label="Decrypted Data")
         
         with gr.TabItem("ℹ️ Info & Key Generation", id=2):
+            # ... UI layout is the same as your working version ...
             gr.Markdown("## Ecosystem Architecture")
-            gr.Markdown(f"This dashboard uses a public [configuration file]({CREATOR_ENDPOINTS_JSON_URL}) to dynamically discover and interact with live services. It demonstrates a secure, decoupled workflow.")
+            gr.Markdown(f"This dashboard uses a public [configuration file]({ENDPOINTS_JSON_URL}) to dynamically discover and interact with live services. It demonstrates a secure, decoupled workflow.")
             with gr.Accordion("🔑 RSA Key Pair Generator", open=False):
                 gr.Markdown("Create a new key pair. In a real scenario, you would add the **Public Key** and the server's **API Endpoint URL** to the `endpoints.json` configuration file, and set the **Private Key** as a secret variable in the corresponding server space.")
                 with gr.Row():
@@ -218,7 +211,11 @@ 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():
-        dropdown_update, status_update, state_update = get_server_list()
+        # The generator must be iterated to get its 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
         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])