Update client/app.py

client/app.py

@@ -1,9 +1,3 @@
-# The client code is excellent as-is. The only change needed is to make it
-# read the endpoints from a local file for the turnkey local demo.
-# Find this line:
-# CREATOR_ENDPOINTS_JSON_URL = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-Creator/raw/main/endpoints.json"
-# And replace the logic to be local-first.
-
 import gradio as gr
 from PIL import Image, ImageDraw, ImageFont, ImageOps
 import base64
@@ -15,6 +9,7 @@ import requests
 import struct
 import numpy as np
 from cryptography.hazmat.primitives import serialization
+# Import 'rsa' for key generation
 from cryptography.hazmat.primitives.asymmetric import rsa, padding
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from cryptography.hazmat.primitives import hashes
@@ -25,15 +20,14 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(level
 logger = logging.getLogger(__name__)
 
 # --- Constants ---
-# Use a local file for easy testing. In a real-world scenario, this could be a URL.
 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 (Your original code is perfect here) ---
-# ...
+# --- Helper Functions for Image Handling ---
+
 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."""
     try:
@@ -43,16 +37,12 @@ def resize_and_crop(img: Image.Image, size: int = IMAGE_SIZE) -> Image.Image:
         return img.resize((size, size), Image.Resampling.LANCZOS)
 
 def prepare_base_image(uploaded_image: Image.Image | None, progress) -> Image.Image:
-    """
-    Provides a base image using the fallback logic, updating a Gradio progress object.
-    """
+    """Provides a base image using a fallback logic, updating a Gradio progress object."""
     if uploaded_image:
         progress(0, desc="✅ Using uploaded image...")
-        logger.info("Using user-uploaded image.")
         return resize_and_crop(uploaded_image)
     try:
-        progress(0, desc="⏳ No image uploaded. Fetching default background...")
-        logger.info(f"Fetching default image from URL: {DEFAULT_IMAGE_URL}")
+        progress(0, desc="⏳ Fetching default background...")
         response = requests.get(DEFAULT_IMAGE_URL, timeout=15)
         response.raise_for_status()
         img = Image.open(io.BytesIO(response.content)).convert("RGB")
@@ -61,23 +51,33 @@ def prepare_base_image(uploaded_image: Image.Image | None, progress) -> Image.Im
     except Exception as e:
         logger.warning(f"Could not fetch default image: {e}. Falling back to AI generation.")
     try:
-        progress(0, desc=f"⏳ Generating new image with {T2I_MODEL}...")
-        logger.info(f"Generating a new image using model: {T2I_MODEL}")
+        progress(0, desc=f"⏳ Generating image with {T2I_MODEL}...")
         client = InferenceClient()
         image_bytes = client.text_to_image(T2I_PROMPT, model=T2I_MODEL)
         img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-        progress(0, desc="✅ New image generated successfully.")
+        progress(0, desc="✅ New image generated.")
         return resize_and_crop(img)
     except Exception as e:
         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 (Your original code is perfect here) ---
-# ...
+# --- Core Cryptography and Image Creation ---
+
 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')
+    """Generates a new 2048-bit RSA private and public key pair."""
+    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(
@@ -86,30 +86,32 @@ def create_encrypted_image(
     base_image: Image.Image,
     overlay_option: 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.")
+    """Encrypts data and embeds it into an image using LSB steganography."""
+    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 in secret data.")
+        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))
+    
+    # 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 overlays
+    # Add Stylish Overlays
+    # (Your existing overlay code is great and remains unchanged)
     draw = ImageDraw.Draw(img, "RGBA")
     try:
         font_bold = ImageFont.truetype("DejaVuSans-Bold.ttf", 30)
@@ -118,93 +120,65 @@ def create_encrypted_image(
         font_bold = ImageFont.load_default(size=28)
         font_regular = ImageFont.load_default(size=14)
 
-    # ... (the rest of your styling logic is great)
-    overlay_color = (15, 23, 42, 190)
-    title_color = (226, 232, 240)
-    key_color = (148, 163, 184)
-    value_color = (241, 245, 249)
+    overlay_color = (15, 23, 42, 190); title_color = (226, 232, 240)
+    key_color = (148, 163, 184); value_color = (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":
-        box_padding = 15
-        line_spacing = 6
-        text_start_x = 35
-
-        if overlay_option == "Keys Only":
-            lines = list(data_dict.keys())
-        else:
-            lines = [f"{key}: {value}" for key, value in data_dict.items()]
-
+        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]
-        total_text_height = sum(line_heights) + (len(lines) - 1) * line_spacing
-        box_height = total_text_height + (box_padding * 2)
+        total_text_height = sum(line_heights) + (len(lines) - 1) * 6
+        box_height = total_text_height + 30
         box_y0 = height - box_height - 20
-
         draw.rectangle([20, box_y0, width - 20, height - 20], fill=overlay_color)
-        current_y = box_y0 + box_padding
-
-        if overlay_option == "Keys Only":
-            for i, key in enumerate(data_dict.keys()):
-                draw.text((text_start_x, current_y), key, fill=key_color, font=font_regular)
-                current_y += line_heights[i] + line_spacing
-        elif overlay_option == "Keys and Values":
-            for i, (key, value) in enumerate(data_dict.items()):
-                key_text = f"{key}:"
-                draw.text((text_start_x, current_y), key_text, fill=key_color, font=font_regular)
-                key_bbox = draw.textbbox((text_start_x, current_y), key_text, font=font_regular)
+        current_y = box_y0 + 15
+        for i, (key, value) in enumerate(data_dict.items()):
+            if overlay_option == "Keys Only":
+                draw.text((35, current_y), key, fill=key_color, font=font_regular)
+            else:
+                key_text = f"{key}:"; draw.text((35, current_y), key_text, fill=key_color, font=font_regular)
+                key_bbox = draw.textbbox((35, current_y), key_text, font=font_regular)
                 draw.text((key_bbox[2] + 8, current_y), str(value), fill=value_color, font=font_regular)
-                current_y += line_heights[i] + line_spacing
+            current_y += line_heights[i] + 6
 
-    # --- Steganography ---
+    # Steganography Payload: [Total Crypto Payload Length (4B)] [Crypto Payload]
     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 size: {pixel_data.size // 8} bytes.")
+        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 ---
+
 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 server configurations found in the file.")
+        if not valid_endpoints: raise ValueError("No valid servers found in config file.")
         endpoint_names = [e['name'] for e in valid_endpoints]
-        status = f"✅ Success! Found {len(endpoint_names)} valid servers."
+        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
     except Exception as e:
         status = f"❌ Error loading configuration: {e}"
-        logger.error(status)
         yield gr.Dropdown(choices=[], value=None, label="Error fetching servers"), status, []
 
-
 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.")
+    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}'. Please refresh the server list.")
+    if not public_key: raise gr.Error(f"Could not find public key for '{service_name}'.")
     try:
         base_image = prepare_base_image(uploaded_image, progress)
         progress(0.5, desc="Encrypting and embedding data...")
-        # Note: The image created here is slightly different from the server example.
-        # This client code uses a different struct packing (`>I` for payload length),
-        # which is correct. The server code will need to match this.
-        # Let's adjust the create_encrypted_image to match server expectations
-        # No, wait, the server code *does* use >I for the *AES key length*, not the *total payload*.
-        # Let's correct the client to be fully compatible.
-
         created_image = create_encrypted_image(secret_data, public_key, base_image, overlay_option)
         return created_image, f"✅ Success! Image created for '{service_name}'.", gr.Tabs(selected=1)
     except Exception as e:
@@ -215,50 +189,43 @@ 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 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 or not endpoint_details.get('link'): raise gr.Error(f"Config Error for '{service_name}'.")
-
+    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')
-    status = f"Connecting to remote server: {server_url}"
-    yield None, status
-
+    api_name = endpoint_details.get('api_endpoint', '/run/keylock-auth-decoder') # Default API name
+    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)
-        # Using the new payload structure for Gradio > 4.2.0
         result = client.predict(image_base64_string=b64_string, api_name=api_name)
-
-        # The result from gradio_client might be the path to a file if it's JSON
+        
+        # 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)
-        else:
-            decrypted_data = result # Assume it's already a dict
-
+             with open(result) 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 calling server with gradio_client: {e}", exc_info=True)
         yield None, f"❌ Error calling server API: {e}"
 
-# --- Gradio UI (Your original code is excellent) ---
-# ... (The rest of your UI code goes here)
-with gr.Blocks(theme="soft", title="KeyLock Operations Dashboard") as demo:
+# --- 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([])
 
     gr.Markdown("# 🔑 KeyLock Operations Dashboard")
-    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.")
+    gr.Markdown("A centralized dashboard to create and test encrypted images against live KeyLock servers.")
 
     with gr.Tabs() as tabs:
         with gr.TabItem("① Create KeyLock", id=0):
-            # ... UI components for creating
-            with gr.Row():
+            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)
-                    creator_secret_input = gr.Textbox(lines=5, label="Secret Data to Encrypt", placeholder="API_KEY: sk-123...\nUSER: demo-user")
-                    creator_base_image_input = gr.Image(label="Optional Base Image (600x600)", type="pil", sources=["upload"])
+                    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_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):
@@ -266,55 +233,63 @@ with gr.Blocks(theme="soft", title="KeyLock Operations Dashboard") as demo:
                     creator_image_output = gr.Image(label="Generated Encrypted Image", type="pil", show_download_button=True, height=600)
 
         with gr.TabItem("② Send KeyLock", id=1):
-             # ... UI components for sending
-            with gr.Row():
+            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 Here", sources=["upload", "clipboard"])
+                    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")
-
+                    client_json_output = gr.JSON(label="Decrypted Data from Server")
+        
         with gr.TabItem("ℹ️ Info & Key Generation", id=2):
-             # ... UI components for info
-            with gr.Accordion("🔑 RSA Key Pair Generator", open=False):
-                output_public_key = gr.Textbox(lines=10, label="Generated Public Key", interactive=False, show_copy_button=True)
-                output_private_key = gr.Textbox(lines=10, label="Generated Private Key", interactive=False, show_copy_button=True)
-                gen_keys_button = gr.Button("⚙️ Generate New 2048-bit Key Pair")
+            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.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)
 
     # --- Event Handlers ---
+    def sync_dropdowns(x): return x
+    
     def refresh_and_update_all():
-        for dropdown_update, status_update, state_update in get_server_list():
-            pass # Exhaust the generator
-        # Sync the dropdowns
-        creator_service_dropdown.update(**dropdown_update.model_dump())
-        send_service_dropdown.update(**dropdown_update.model_dump())
-        return dropdown_update, status_update, state_update
+        # 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=get_server_list,
-        outputs=[creator_service_dropdown, creator_status, endpoints_state]
-    ).then(
-        lambda x: x,
-        inputs=creator_service_dropdown,
-        outputs=send_service_dropdown
+    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]
     )
-
-    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]
     )
+    
     client_button.click(
-        fn=send_keylock_wrapper,
-        inputs=[send_service_dropdown, client_image_input, endpoints_state],
+        fn=send_keylock_wrapper, 
+        inputs=[send_service_dropdown, client_image_input, endpoints_state], 
         outputs=[client_json_output, client_status]
     )
     
-    # When a new image is created, also update the sender tab's image
-    creator_image_output.change(lambda x: x, inputs=creator_image_output, outputs=client_image_input)
+    # 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__":
     demo.launch()