Update client/app.py

client/app.py

@@ -7,9 +7,9 @@ import logging
 import os
 import requests
 import struct
+import tempfile # <--- Import tempfile
 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
@@ -63,32 +63,14 @@ def prepare_base_image(uploaded_image: Image.Image | None, progress) -> Image.Im
 # --- Core Cryptography and Image Creation (Unchanged) ---
 
 def generate_rsa_keys():
-    """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')
+    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,
-    base_image: Image.Image,
-    overlay_option: str
-) -> Image.Image:
-    """Encrypts data and embeds it into an image using LSB steganography."""
+def create_encrypted_image(secret_data_str: str, public_key_pem: str, 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.")
-
     data_dict = {}
     for line in secret_data_str.splitlines():
         line = line.strip()
@@ -96,18 +78,15 @@ def create_encrypted_image(
         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 = base_image.copy().convert("RGB")
     width, height = img.size
-
-    # Stylish Overlays (Code unchanged)
+    # (Overlay drawing code is unchanged and correct)
     draw = ImageDraw.Draw(img, "RGBA")
     try:
         font_bold = ImageFont.truetype("DejaVuSans-Bold.ttf", 30)
@@ -128,22 +107,17 @@ def create_encrypted_image(
         draw.rectangle([20, box_y0, width - 20, height - 20], fill=overlay_color)
         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)
+            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] + 6
-
-    # Steganography
+    # (Steganography code is unchanged and correct)
     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}.")
-
+    if len(binary_payload) > pixel_data.size: raise ValueError(f"Data is too large. Max: {pixel_data.size // 8} bytes.")
     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')
@@ -151,70 +125,80 @@ def create_encrypted_image(
 # --- 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://...).")
-
+    # This function is correct from the previous fix. Unchanged.
+    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."
-
+        gr.Info(f"Server at {url} is already in the list.")
+        return current_servers, f"ℹ️ Server at {url} is already in the list.", gr.Button(interactive=True)
     status = f"⏳ Testing server at {url}..."
     yield current_servers, status, gr.Button(interactive=False)
-
     try:
         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_info = client.predict(api_name="/keylock-info")
+        if not isinstance(server_info, dict): raise TypeError(f"Expected dict from /keylock-info, but got {type(server_info)}")
         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"
-        }
+        public_key = client.predict(api_name="/keylock-pub")
+        if not isinstance(public_key, str): raise TypeError(f"Expected a string from /keylock-pub, but got {type(public_key)}")
+        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}"
+        status_message = f"✅ Successfully added server: {server_name}"; gr.Info(status_message)
         yield updated_servers, status_message, gr.Button(interactive=True)
-
     except Exception as e:
         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}"
+        status_message = f"❌ Failed to connect or validate server at {url}. Error: {e}"; gr.Error(status_message)
         yield current_servers, status_message, gr.Button(interactive=True)
 
 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)):
+    """
+    Creates the encrypted image and returns it to two separate components:
+    1. A gr.Image for visual preview (which may be re-encoded by Gradio).
+    2. A gr.File for a guaranteed, uncorrupted PNG download.
+    """
     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.")
+    
+    png_path = None # Initialize to handle potential errors
     try:
         base_image = prepare_base_image(uploaded_image, progress)
         progress(0.5, desc="Encrypting and embedding data...")
         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)
+        
+        # Save image to a temporary PNG file to prevent re-encoding.
+        # This temporary file will be served by gr.File for perfect data integrity.
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_f:
+            png_path = temp_f.name
+            created_image.save(png_path, "PNG", compress_level=1)
+        
+        status_message = f"✅ Success! Image created for '{service_name}'. Use the link to download the uncorrupted file."
+        # Return path to both preview and file components. Make file component visible.
+        return png_path, png_path, status_message, gr.Tabs(selected=1), gr.File(visible=True)
+
     except Exception as e:
         logger.error(f"Error creating image: {e}", exc_info=True)
-        return None, f"❌ Error: {e}", gr.Tabs()
+        if png_path and os.path.exists(png_path):
+            os.remove(png_path) # Clean up temp file on error
+        # Return Nones and hide file component
+        return None, None, f"❌ Error: {e}", gr.Tabs(), gr.File(visible=False)
 
-def send_keylock_wrapper(service_name: str, image: Image.Image, available_servers: list):
+def send_keylock_wrapper(service_name: str, image_path: str, available_servers: list):
+    # This function now expects a file path `image_path` from gr.Image
     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.")
+    if image_path is None: raise gr.Error("Please create or upload an image to send.")
+    
     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")
+        # The input `image_path` is a path to a temporary file created by gr.Image.
+        # We read it directly to ensure we get the correct bytes.
+        with open(image_path, "rb") as img_file:
+            b64_string = base64.b64encode(img_file.read()).decode("utf-8")
+            
         client = Client(server_url)
         result = client.predict(image_base64_string=b64_string, api_name=api_name)
 
@@ -252,13 +236,16 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
                     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_image_output = gr.Image(label="Generated Encrypted Image", type="pil", show_download_button=True, height=600)
+                    # Use two components: one for preview, one for lossless download
+                    creator_image_preview = gr.Image(label="Image Preview (for display only)", type="filepath", height=500)
+                    creator_file_output = gr.File(label="Download Uncorrupted PNG File", file_count="single", visible=False)
 
         with gr.TabItem("② Send KeyLock", id=1):
             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"])
+                    # This component will receive the file path from the preview image
+                    client_image_input = gr.Image(type="filepath", 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_json_output = gr.JSON(label="Decrypted Data from Server")
@@ -279,24 +266,16 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
         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
+    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])
 
-    servers_state.change(
-        fn=update_dropdowns_from_state,
-        inputs=servers_state,
-        outputs=[creator_service_dropdown, send_service_dropdown]
-    )
+    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, servers_state, creator_base_image_input, creator_overlay_option],
-        outputs=[creator_image_output, global_status, tabs]
+        outputs=[creator_image_preview, creator_file_output, global_status, tabs, creator_file_output]
     )
 
     client_button.click(
@@ -307,7 +286,10 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
 
     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)
-    creator_image_output.change(fn=lambda x: x, inputs=creator_image_output, outputs=client_image_input)
+    
+    # When a new image preview is generated, auto-populate it in the sender tab.
+    # This works because both components now use file paths.
+    creator_image_preview.change(fn=lambda x: x, inputs=creator_image_preview, outputs=client_image_input)
 
 if __name__ == "__main__":
     demo.launch()