Update app.py

app.py

@@ -18,9 +18,8 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(level
 logger = logging.getLogger(__name__)
 
 # ==============================================================================
-#  CONFIGURATION: URL to the public list of services
+#  CONFIGURATION
 # ==============================================================================
-# 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"
 
 # ==============================================================================
@@ -29,13 +28,8 @@ 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:
@@ -43,22 +37,15 @@ 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 = {}
-    for line in secret_data_str.splitlines():
-        if not line.strip() or line.strip().startswith('#'): continue
-        parts = line.split(':', 1) if ':' in line else line.split('=', 1)
-        if len(parts) != 2: continue
-        data_dict[parts[0].strip()] = parts[1].strip().strip("'\"")
+    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}
     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_with_tag = AESGCM(aes_key).encrypt(nonce, json_bytes, None)
-    rsa_encrypted_aes_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_aes_key)) + rsa_encrypted_aes_key + nonce + ciphertext_with_tag
+    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)
@@ -67,9 +54,10 @@ def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.I
     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(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])
+    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')
 
@@ -79,18 +67,19 @@ def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.I
 
 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()
+        # 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."
-        yield gr.Dropdown(choices=endpoint_names, value=endpoint_names[0] if endpoint_names else None, label="Target Server"), status, endpoints
+        return 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, []
+        return 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."""
@@ -108,13 +97,14 @@ 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.")
     
-    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
+    # --- 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)
+    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)
     
     try:
         with io.BytesIO() as buffer:
@@ -153,7 +143,7 @@ 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.")
+    gr.Markdown("A centralized dashboard to manage and demonstrate the entire 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("① Create KeyLock", id=0):
@@ -163,7 +153,7 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
                 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):
@@ -172,12 +162,12 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
 
         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.")
+            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"):
                 with gr.Column(scale=1):
-                    gr.Markdown("### Select Server to Call")
+                    gr.Markdown("### Configuration")
                     send_service_dropdown = gr.Dropdown(label="Target Server", interactive=True, info="Select the API server to send the image to.")
-                    gr.Markdown("### Upload Image")
+                    gr.Markdown("### Image to Send")
                     client_image_input = gr.Image(type="pil", label="Upload or Drag Encrypted Image Here", sources=["upload", "clipboard"])
                     client_button = gr.Button("🔓 Decrypt via Remote Server", variant="primary")
                 with gr.Column(scale=1):
@@ -185,11 +175,11 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
                     client_status = gr.Textbox(label="Status", interactive=False, lines=2)
                     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.TabItem("ℹ️ Info & Key Generation", id=2):
+            gr.Markdown("## Ecosystem Architecture")
+            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 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.")
+                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():
                     with gr.Column():
                         output_public_key = gr.Textbox(lines=10, label="Generated Public Key", interactive=False, show_copy_button=True)
@@ -201,9 +191,8 @@ 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():
-        *_, last_yield = get_server_list()
-        dropdown_update, status_update, state_update = last_yield
-        # Update both dropdowns simultaneously
+        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])