Update app.py

app.py

@@ -1,146 +1,3 @@
-import gradio as gr
-from PIL import Image, ImageDraw, ImageFont
-import base64
-import io
-import json
-import logging
-import os
-import requests
-import struct
-import numpy as np
-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
-# ==============================================================================
-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}"
-
-# ==============================================================================
-#  LOCAL LOGIC (Key and Image Generation)
-# ==============================================================================
-def generate_rsa_keys():
-    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')
-    return private_pem, public_pem
-
-def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.Image:
-    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():
-        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: 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
-# ==============================================================================
-
-def get_server_list():
-    """Fetches and validates 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()
-        
-        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:
-        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.")
-    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.")
-    
-    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}"
-    yield None, status
-    
-    try:
-        with io.BytesIO() as buffer:
-            image.save(buffer, format="PNG")
-            b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
-        
-        payload = {"data": [b64_string]}
-        headers = {"Content-Type": "application/json"}
-        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:
-                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:
-        yield None, f"❌ Error calling server API: {e}"
-
 # ==============================================================================
 #  GRADIO DASHBOARD INTERFACE
 # ==============================================================================
@@ -154,6 +11,7 @@ theme = gr.themes.Base(
 )
 
 with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
+    # This State component will hold the list of dicts fetched from the JSON config
     endpoints_state = gr.State([])
 
     gr.Markdown("# 🔑 KeyLock Operations Dashboard")
@@ -161,14 +19,13 @@ 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):
@@ -176,7 +33,6 @@ 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"):
@@ -192,9 +48,14 @@ 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.Row():
+                 with gr.Column():
+                    gr.Markdown(f"### 🏭 Auth Creator Service\n- **Space:** [{CREATOR_SPACE_ID}]({CREATOR_URL})\n- **Role:** Provides an API to encrypt data for various targets defined in its `endpoints.json` file.\n- **Source Code:** [app.py]({CREATOR_APP_PY_URL})")
+                 with gr.Column():
+                    gr.Markdown(f"### 📡 Decoder Server\n- **Space:** [{SERVER_SPACE_ID}]({SERVER_URL})\n- **Role:** The trusted authority. It holds a secret private key and provides a secure API to decrypt images.\n- **Source Code:** [app.py]({SERVER_APP_PY_URL})")
+            
             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():
@@ -209,8 +70,10 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
     
     def refresh_and_update_all():
         # The generator must be iterated to get its final values
+        # The 'pass' allows the loop to complete, and we use the final yielded values.
         for dropdown_update, status_update, state_update in get_server_list():
             pass
+        # 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])