Update app.py

app.py

@@ -1,26 +1,39 @@
 import gradio as gr
-import json
-import os
-import io
+from gradio_client import Client
+from PIL import Image, ImageDraw, ImageFont
 import base64
-import struct
+import io
+import json
 import logging
+import os
 import requests
-from PIL import Image, ImageDraw, ImageFont
+import tempfile
+import struct
 import numpy as np
-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.hazmat.primitives.ciphers.aead import AESGCM
+from cryptography.hazmat.primitives import hashes
 
 # --- Configure Logging ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 # ==============================================================================
-#  CONFIGURATION
+#  CONFIGURATION: IDs AND URLs OF THE REMOTE SERVICES
 # ==============================================================================
-ENDPOINTS_JSON_URL = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-Creator/raw/main/endpoints.json"
+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}"
+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)
@@ -28,8 +41,13 @@ ENDPOINTS_JSON_URL = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-
 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:
@@ -37,7 +55,20 @@ def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.I
     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.")
     
-    data_dict = {parts[0].strip(): parts[1].strip().strip("'\"") for line in secret_data_str.splitlines() if (line := line.strip()) and not line.startswith('#') and (parts := (line.split(':', 1) if ':' in line else line.split('=', 1))) and len(parts) == 2}
+    # --- 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
+        
+        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 not data_dict: raise ValueError("No valid key-value pairs found.")
     
     json_bytes = json.dumps(data_dict).encode('utf-8')
@@ -62,24 +93,22 @@ def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.I
     return Image.fromarray(stego_pixels, 'RGB')
 
 # ==============================================================================
-#  UI HELPER & REMOTE API CALL LOGIC
+#  UI HELPER & REMOTE API CALL LOGIC (Your working versions)
 # ==============================================================================
-
-def get_server_list():
-    """Fetches the list of servers from the public JSON file."""
+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
     try:
-        response = requests.get(ENDPOINTS_JSON_URL, timeout=10)
+        response = requests.get(CREATOR_ENDPOINTS_JSON_URL, timeout=10)
         response.raise_for_status()
         endpoints = response.json()
-        # Validate the new structure
-        for e in endpoints:
-            if "api_endpoint" not in e: raise ValueError("Configuration Error: Entry is missing the required 'api_endpoint' field.")
         endpoint_names = [e['name'] for e in endpoints]
-        status = f"✅ Success! Found {len(endpoint_names)} servers."
-        return gr.Dropdown(choices=endpoint_names, value=endpoint_names[0] if endpoint_names else None, label="Target Server"), status, 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
     except Exception as e:
-        status = f"❌ Error fetching configuration: {e}"
-        return gr.Dropdown(choices=[], value=None, label="Error fetching servers"), status, []
+        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, []
 
 def create_keylock_wrapper(service_name: str, secret_data: str, available_endpoints: list):
     """UI wrapper for creating an image for a selected service."""
@@ -97,14 +126,12 @@ def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoi
     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.")
     
-    # --- THIS IS THE FIX ---
-    # Look up the correct, full API endpoint from the configuration. No more guesswork!
-    api_endpoint = next((e['api_endpoint'] for e in available_endpoints if e['name'] == service_name), None)
+    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}'.")
 
     status = f"Connecting to endpoint: {api_endpoint}"
-    logger.info(status)
+    yield None, status
     
     try:
         with io.BytesIO() as buffer:
@@ -128,7 +155,7 @@ def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoi
         return None, f"❌ Error calling server API: {e}"
 
 # ==============================================================================
-#  GRADIO DASHBOARD INTERFACE
+#  GRADIO DASHBOARD INTERFACE (Your working layout)
 # ==============================================================================
 theme = gr.themes.Base(
     primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky, neutral_hue=gr.themes.colors.slate,
@@ -192,7 +219,6 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
     
     def refresh_and_update_all():
         dropdown_update, status_update, state_update = get_server_list()
-        # Update both dropdowns simultaneously from the same source of truth
         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])