Update client/app.py

client/app.py

@@ -20,13 +20,12 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(level
 logger = logging.getLogger(__name__)
 
 # --- Constants ---
-ENDPOINTS_CONFIG_PATH = "endpoints.json"
 DEFAULT_IMAGE_URL = "https://images.unsplash.com/photo-1506318137071-a8e063b4bec0?q=80&w=1200&auto=format&fit=crop"
 IMAGE_SIZE = 600
 T2I_MODEL = "sd-community/sdxl-lightning"
 T2I_PROMPT = "A stunning view of a distant galaxy, nebulae, and constellations, digital art, vibrant colors, cinematic lighting, 8k, masterpiece."
 
-# --- Helper Functions for Image Handling ---
+# --- Helper Functions for Image Handling (Unchanged) ---
 
 def resize_and_crop(img: Image.Image, size: int = IMAGE_SIZE) -> Image.Image:
     """Resizes an image to fit within a square of `size` and then center-crops it."""
@@ -61,7 +60,7 @@ def prepare_base_image(uploaded_image: Image.Image | None, progress) -> Image.Im
         logger.error(f"Fatal: All image sources failed. Text-to-image failed with: {e}")
         raise gr.Error(f"Failed to obtain a base image. AI generation error: {e}")
 
-# --- Core Cryptography and Image Creation ---
+# --- Core Cryptography and Image Creation (Unchanged) ---
 
 def generate_rsa_keys():
     """Generates a new 2048-bit RSA private and public key pair."""
@@ -103,15 +102,12 @@ def create_encrypted_image(
     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))
-    
-    # Payload format: [RSA Enc Key Len (4B)] [RSA Enc Key (var)] [Nonce (12B)] [Ciphertext (var)]
+
     encrypted_payload = struct.pack('>I', len(rsa_encrypted_key)) + rsa_encrypted_key + nonce + ciphertext
-    
     img = base_image.copy().convert("RGB")
     width, height = img.size
 
-    # Add Stylish Overlays
-    # (Your existing overlay code is great and remains unchanged)
+    # Stylish Overlays (Code unchanged)
     draw = ImageDraw.Draw(img, "RGBA")
     try:
         font_bold = ImageFont.truetype("DejaVuSans-Bold.ttf", 30)
@@ -119,13 +115,10 @@ def create_encrypted_image(
     except IOError:
         font_bold = ImageFont.load_default(size=28)
         font_regular = ImageFont.load_default(size=14)
-
-    overlay_color = (15, 23, 42, 190); title_color = (226, 232, 240)
-    key_color = (148, 163, 184); value_color = (241, 245, 249)
-
+    overlay_color, title_color = (15, 23, 42, 190), (226, 232, 240)
+    key_color, value_color = (148, 163, 184), (241, 245, 249)
     draw.rectangle([0, 20, width, 80], fill=overlay_color)
     draw.text((width / 2, 50), "KeyLock Secure Data", fill=title_color, font=font_bold, anchor="ms")
-
     if overlay_option != "None":
         lines = list(data_dict.keys()) if overlay_option == "Keys Only" else [f"{k}: {v}" for k, v in data_dict.items()]
         line_heights = [draw.textbbox((0, 0), line, font=font_regular)[3] for line in lines]
@@ -143,152 +136,177 @@ def create_encrypted_image(
                 draw.text((key_bbox[2] + 8, current_y), str(value), fill=value_color, font=font_regular)
             current_y += line_heights[i] + 6
 
-    # Steganography Payload: [Total Crypto Payload Length (4B)] [Crypto Payload]
+    # Steganography
     pixel_data = np.array(img.convert("RGB")).ravel()
     header = struct.pack('>I', len(encrypted_payload))
     binary_payload = ''.join(format(b, '08b') for b in header + encrypted_payload)
-    
+
     if len(binary_payload) > pixel_data.size:
         raise ValueError(f"Data is too large for the image. Max bytes: {pixel_data.size // 8}.")
-    
+
     pixel_data[:len(binary_payload)] = (pixel_data[:len(binary_payload)] & 0xFE) | np.array(list(binary_payload), dtype=np.uint8)
     stego_pixels = pixel_data.reshape((height, width, 3))
     return Image.fromarray(stego_pixels, 'RGB')
 
-# --- Gradio Wrapper Functions ---
+# --- New Server Management and Gradio Wrappers ---
+
+def add_server(url: str, current_servers: list):
+    """Tests a server URL, fetches its public key and info, and adds it to the list."""
+    if not url or not url.startswith(('http://', 'https://')):
+        raise gr.Error("Please enter a valid server URL (e.g., https://...).")
+
+    url = url.strip().rstrip('/')
+    if any(s['url'] == url for s in current_servers):
+        return current_servers, f"ℹ️ Server at {url} is already in the list."
+
+    status = f"⏳ Testing server at {url}..."
+    yield current_servers, status, gr.Button(interactive=False)
 
-def get_server_list():
-    status = f"Fetching server list from '{ENDPOINTS_CONFIG_PATH}'..."
-    yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, []
     try:
-        with open(ENDPOINTS_CONFIG_PATH, "r") as f:
-            all_entries = json.load(f)
-        valid_endpoints = [e for e in all_entries if isinstance(e, dict) and "name" in e and "public_key" in e and "link" in e]
-        if not valid_endpoints: raise ValueError("No valid servers found in config file.")
-        endpoint_names = [e['name'] for e in valid_endpoints]
-        status = f"✅ 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
+        client = Client(url, verbose=False)
+
+        info_result_path = client.predict(api_name="/keylock-info")
+        with open(info_result_path, 'r', encoding='utf-8') as f: server_info = json.load(f)
+        server_name = server_info.get("name", url)
+
+        pubkey_result_path = client.predict(api_name="/keylock-pub")
+        with open(pubkey_result_path, 'r', encoding='utf-8') as f: public_key = f.read()
+
+        if not public_key.strip().startswith("-----BEGIN PUBLIC KEY-----"):
+            raise ValueError("Received invalid public key format from server.")
+
+        new_server = {
+            "name": f"{server_name} ({url.split('//')[1].split('/')[0]})",
+            "url": url,
+            "public_key": public_key,
+            "api_endpoint": "/keylock-server"
+        }
+        updated_servers = current_servers + [new_server]
+        status_message = f"✅ Successfully added server: {server_name}"
+        yield updated_servers, status_message, gr.Button(interactive=True)
+
     except Exception as e:
-        status = f"❌ Error loading configuration: {e}"
-        yield gr.Dropdown(choices=[], value=None, label="Error fetching servers"), status, []
+        logger.error(f"Failed to add server at {url}: {e}", exc_info=True)
+        status_message = f"❌ Failed to connect or validate server at {url}. Error: {e}"
+        yield current_servers, status_message, gr.Button(interactive=True)
 
-def create_keylock_wrapper(service_name: str, secret_data: str, available_endpoints: list, uploaded_image: Image.Image | None, overlay_option: str, progress=gr.Progress(track_tqdm=True)):
-    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}'.")
+def create_keylock_wrapper(service_name: str, secret_data: str, available_servers: list, uploaded_image: Image.Image | None, overlay_option: str, progress=gr.Progress(track_tqdm=True)):
+    if not service_name: raise gr.Error("Please add and select a target server.")
+    server = next((s for s in available_servers if s['name'] == service_name), None)
+    if not server: raise gr.Error(f"Could not find config for '{service_name}'. Please re-add it.")
     try:
         base_image = prepare_base_image(uploaded_image, progress)
         progress(0.5, desc="Encrypting and embedding data...")
-        created_image = create_encrypted_image(secret_data, public_key, base_image, overlay_option)
+        created_image = create_encrypted_image(secret_data, server['public_key'], base_image, overlay_option)
         return created_image, f"✅ Success! Image created for '{service_name}'.", gr.Tabs(selected=1)
     except Exception as e:
         logger.error(f"Error creating image: {e}", exc_info=True)
         return None, f"❌ Error: {e}", gr.Tabs()
 
-def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoints: list):
+def send_keylock_wrapper(service_name: str, image: Image.Image, available_servers: list):
     if not service_name: raise gr.Error("Please select a target server.")
     if image is None: raise gr.Error("Please create or upload an image to send.")
-    endpoint_details = next((e for e in available_endpoints if e['name'] == service_name), None)
-    if not endpoint_details: raise gr.Error(f"Config Error for '{service_name}'.")
-    server_url = endpoint_details['link']
-    api_name = endpoint_details.get('api_endpoint', '/run/keylock-auth-decoder') # Default API name
-    yield None, f"Connecting to remote server: {server_url}"
+    server = next((s for s in available_servers if s['name'] == service_name), None)
+    if not server: raise gr.Error(f"Config Error for '{service_name}'. Please re-add it.")
+
+    server_url, api_name = server['url'], server['api_endpoint']
+    yield None, f"⏳ Connecting to remote server: {server_url}"
     try:
         with io.BytesIO() as buffer:
             image.save(buffer, format="PNG")
             b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
         client = Client(server_url)
         result = client.predict(image_base64_string=b64_string, api_name=api_name)
-        
-        # Gradio_client can return a filepath for JSON, so we handle that case.
+
         if isinstance(result, str) and os.path.exists(result):
-             with open(result) as f: decrypted_data = json.load(f)
+             with open(result, 'r', encoding='utf-8') as f: decrypted_data = json.load(f)
         else: decrypted_data = result
-            
+
         yield decrypted_data, "✅ Success! Data decrypted by remote server."
     except Exception as e:
+        logger.error(f"Error sending to server: {e}", exc_info=True)
         yield None, f"❌ Error calling server API: {e}"
 
 # --- Gradio UI ---
 theme = gr.themes.Soft(primary_hue="blue", secondary_hue="sky", neutral_hue="slate")
 
 with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
-    endpoints_state = gr.State([])
+    servers_state = gr.State([])
 
     gr.Markdown("# 🔑 KeyLock Operations Dashboard")
-    gr.Markdown("A centralized dashboard to create and test encrypted images against live KeyLock servers.")
+    gr.Markdown("A client to discover KeyLock servers, create encrypted images, and test decryption against live APIs.")
+
+    with gr.Accordion("🔗 Add a KeyLock Server", open=True):
+        with gr.Row():
+            server_url_input = gr.Textbox(label="Server Base URL", placeholder="https://your-hf-space-name.hf.space")
+            add_server_button = gr.Button("Test & Add Server", variant="secondary")
+    global_status = gr.Textbox(label="Status", interactive=False, lines=1, value="Enter a server URL to begin.")
 
     with gr.Tabs() as tabs:
         with gr.TabItem("① Create KeyLock", id=0):
-            gr.Markdown("## Step 1: Create an Encrypted Image (Local)")
             with gr.Row(variant="panel"):
                 with gr.Column(scale=2):
-                    creator_service_dropdown = gr.Dropdown(label="Target Server", interactive=True, info="Select the API server to encrypt data for.")
-                    creator_secret_input = gr.Textbox(lines=5, label="Secret Data", placeholder="API_KEY: sk-123...\nUSER: demo-user")
+                    creator_service_dropdown = gr.Dropdown(label="Target Server", interactive=True, info="Select a server to encrypt data for.")
+                    creator_secret_input = gr.Textbox(lines=5, label="Secret Data (key:value format)", placeholder="API_KEY: sk-123...\nUSER: demo-user")
                     creator_base_image_input = gr.Image(label="Optional Base Image (600x600 recommended)", type="pil", sources=["upload"])
                     creator_overlay_option = gr.Radio(label="Overlay Content", choices=["Keys and Values", "Keys Only", "None"], value="Keys and Values")
                     creator_button = gr.Button("✨ Create Auth Image", variant="primary")
                 with gr.Column(scale=3):
-                    creator_status = gr.Textbox(label="Status", interactive=False, lines=1)
                     creator_image_output = gr.Image(label="Generated Encrypted Image", type="pil", show_download_button=True, height=600)
 
         with gr.TabItem("② Send KeyLock", id=1):
-            gr.Markdown("## Step 2: Decrypt via Live API Call")
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1):
                     send_service_dropdown = gr.Dropdown(label="Target Server", interactive=True)
                     client_image_input = gr.Image(type="pil", label="Upload or Drag Encrypted Image", sources=["upload", "clipboard"])
                     client_button = gr.Button("🔓 Decrypt via Remote Server", variant="primary")
                 with gr.Column(scale=1):
-                    client_status = gr.Textbox(label="Status", interactive=False, lines=2)
                     client_json_output = gr.JSON(label="Decrypted Data from Server")
-        
-        with gr.TabItem("ℹ️ Info & Key Generation", id=2):
-            gr.Markdown("## About the KeyLock System")
-            gr.Markdown("This dashboard uses the `endpoints.json` file to discover servers. Each server holds a secret **private key**. This client uses the corresponding **public key** to encrypt data that only that specific server can decrypt.")
-            
+
+        with gr.TabItem("ℹ️ Key Generation", id=2):
             with gr.Accordion("🔑 RSA Key Pair Generator", open=True):
-                gr.Markdown("Use this utility to generate a new key pair for a new server. Add the **Public Key** to `endpoints.json` and set the **Private Key** as a secret on the server.")
-                gen_keys_button = gr.Button("⚙️ Generate New 2048-bit Key Pair", variant="secondary")
-                with gr.Row():
-                    with gr.Column():
-                        output_public_key = gr.Textbox(lines=10, label="Generated Public Key", interactive=False, show_copy_button=True)
-                    with gr.Column():
-                        output_private_key = gr.Textbox(lines=10, label="Generated Private Key (Keep this secret!)", interactive=False, show_copy_button=True)
+                gr.Markdown("Use this utility to generate a new key pair for a new server. The server administrator must configure the server with the private key.")
+                gen_keys_button = gr.Button("⚙️ Generate New 2048-bit Key Pair")
+                output_private_key = gr.Textbox(lines=10, label="Generated Private Key (KEEP THIS SECRET!)", interactive=False, show_copy_button=True)
+                output_public_key = gr.Textbox(lines=10, label="Generated Public Key", interactive=False, show_copy_button=True)
 
     # --- Event Handlers ---
     def sync_dropdowns(x): return x
-    
-    def refresh_and_update_all():
-        # This generator pattern is for Gradio's streaming output
-        for dropdown_update, status_update, state_update in get_server_list(): pass
-        return dropdown_update, dropdown_update, status_update, state_update
-
-    demo.load(fn=refresh_and_update_all, outputs=[creator_service_dropdown, send_service_dropdown, creator_status, endpoints_state])
-    
-    gen_keys_button.click(
-        fn=generate_rsa_keys, 
-        inputs=None, 
-        outputs=[output_private_key, output_public_key]
+
+    def update_dropdowns_from_state(servers_list):
+        server_names = [s['name'] for s in servers_list] if servers_list else []
+        selected = server_names[-1] if server_names else None
+        dropdown_update = gr.Dropdown(choices=server_names, value=selected, label="Target Server", interactive=True)
+        return dropdown_update, dropdown_update
+
+    add_server_button.click(
+        fn=add_server,
+        inputs=[server_url_input, servers_state],
+        outputs=[servers_state, global_status, add_server_button]
+    ).then(fn=lambda: "", outputs=[server_url_input]) # Clear input after trying
+
+    servers_state.change(
+        fn=update_dropdowns_from_state,
+        inputs=servers_state,
+        outputs=[creator_service_dropdown, send_service_dropdown]
     )
-    
+
+    gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
+
     creator_button.click(
         fn=create_keylock_wrapper,
-        inputs=[creator_service_dropdown, creator_secret_input, endpoints_state, creator_base_image_input, creator_overlay_option],
-        outputs=[creator_image_output, creator_status, tabs]
+        inputs=[creator_service_dropdown, creator_secret_input, servers_state, creator_base_image_input, creator_overlay_option],
+        outputs=[creator_image_output, global_status, tabs]
     )
-    
+
     client_button.click(
-        fn=send_keylock_wrapper, 
-        inputs=[send_service_dropdown, client_image_input, endpoints_state], 
-        outputs=[client_json_output, client_status]
+        fn=send_keylock_wrapper,
+        inputs=[send_service_dropdown, client_image_input, servers_state],
+        outputs=[client_json_output, global_status]
     )
-    
-    # Sync the creator/sender dropdowns when one changes
+
     creator_service_dropdown.change(fn=sync_dropdowns, inputs=creator_service_dropdown, outputs=send_service_dropdown)
     send_service_dropdown.change(fn=sync_dropdowns, inputs=send_service_dropdown, outputs=creator_service_dropdown)
-    
-    # When a new image is created, auto-populate it in the sender tab
     creator_image_output.change(fn=lambda x: x, inputs=creator_image_output, outputs=client_image_input)
 
 if __name__ == "__main__":