Update app.py

app.py

@@ -39,11 +39,9 @@ T2I_PROMPT = "A stunning view of a distant galaxy, nebulae, and constellations,
 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:
-        # ImageOps.fit is perfect for this: it resizes while maintaining aspect ratio and crops from the center.
         return ImageOps.fit(img, (size, size), Image.Resampling.LANCZOS)
     except Exception as e:
         logger.error(f"Failed to resize and crop image: {e}")
-        # Fallback to a simple resize if 'fit' fails
         return img.resize((size, size), Image.Resampling.LANCZOS)
 
 def prepare_base_image(uploaded_image: Image.Image | None, progress) -> Image.Image:
@@ -54,13 +52,10 @@ def prepare_base_image(uploaded_image: Image.Image | None, progress) -> Image.Im
     3. AI-generated image.
     All images are resized and cropped to a standard size.
     """
-    # 1. User-uploaded image
     if uploaded_image:
         progress(0, desc="✅ Using uploaded image...")
         logger.info("Using user-uploaded image.")
         return resize_and_crop(uploaded_image)
-
-    # 2. Default URL image
     try:
         progress(0, desc="⏳ No image uploaded. Fetching default background...")
         logger.info(f"Fetching default image from URL: {DEFAULT_IMAGE_URL}")
@@ -71,8 +66,6 @@ def prepare_base_image(uploaded_image: Image.Image | None, progress) -> Image.Im
         return resize_and_crop(img)
     except Exception as e:
         logger.warning(f"Could not fetch default image: {e}. Falling back to AI generation.")
-
-    # 3. AI-generated image (last resort)
     try:
         progress(0, desc=f"⏳ Generating new image with {T2I_MODEL}...")
         logger.info(f"Generating a new image using model: {T2I_MODEL}")
@@ -97,13 +90,12 @@ def create_encrypted_image(
     secret_data_str: str,
     public_key_pem: str,
     base_image: Image.Image,
-    show_key_overlay: bool
+    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()
@@ -111,79 +103,73 @@ 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 in secret data.")
-        
     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
     
-    # Use the prepared base_image
     img = base_image.copy().convert("RGB")
     width, height = img.size
 
     # --- Add Stylish Overlays ---
-    draw = ImageDraw.Draw(img, "RGBA") # Use RGBA to draw transparent shapes
+    draw = ImageDraw.Draw(img, "RGBA")
     try:
         font_bold = ImageFont.truetype("DejaVuSans-Bold.ttf", 30)
         font_regular = ImageFont.truetype("DejaVuSans.ttf", 15)
-        font_small = ImageFont.truetype("DejaVuSans.ttf", 12)
     except IOError:
         font_bold = ImageFont.load_default(size=28)
         font_regular = ImageFont.load_default(size=14)
-        font_small = ImageFont.load_default(size=11)
 
-    # Colors matching the Gradio theme (slate, sky)
-    overlay_color = (15, 23, 42, 190)  # very dark slate, semi-transparent
-    title_color = (226, 232, 240)      # light slate/sky
-    key_color = (148, 163, 184)        # slate-400 for keys
-    value_color = (241, 245, 249)      # slate-100 for values
+    overlay_color = (15, 23, 42, 190)
+    title_color = (226, 232, 240)
+    key_color = (148, 163, 184)
+    value_color = (241, 245, 249)
 
-    # Main title overlay (top)
     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")
 
-    # Key/Value data overlay (bottom-left)
-    if show_key_overlay:
+    # Data overlay logic based on user's choice
+    if overlay_option != "None":
         box_padding = 15
         line_spacing = 6
         text_start_x = 35
-        
-        # Calculate box height dynamically
-        lines = [f"{key}: {value}" for key, value in data_dict.items()]
-        line_heights = [draw.textbbox((0,0), line, font=font_regular)[3] for line in lines]
+
+        if overlay_option == "Keys Only":
+            lines = list(data_dict.keys())
+        else: # "Keys and Values"
+            lines = [f"{key}: {value}" for key, value 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)
         box_y0 = height - box_height - 20
         
         draw.rectangle([20, box_y0, width - 20, height - 20], fill=overlay_color)
-
         current_y = box_y0 + box_padding
-        for i, (key, value) in enumerate(data_dict.items()):
-            # Draw key and value separately for different colors
-            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)
-            draw.text((key_bbox[2] + 8, current_y), str(value), fill=value_color, font=font_regular)
-            current_y += line_heights[i] + line_spacing
-
-    # --- Steganography: Embed encrypted data into pixels ---
+        
+        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)
+                draw.text((key_bbox[2] + 8, current_y), str(value), fill=value_color, font=font_regular)
+                current_y += line_heights[i] + line_spacing
+
+    # --- Steganography ---
     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(f"Data is too large for the image. Max size: {pixel_data.size // 8} bytes. Your data: ~{len(binary_payload) // 8} bytes.")
-        
+        raise ValueError(f"Data is too large for the image. Max size: {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')
 
 # --- Gradio Wrapper Functions ---
@@ -206,20 +192,16 @@ def get_server_list():
         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, show_keys: bool, progress=gr.Progress(track_tqdm=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}'. Please refresh the server list.")
     try:
-        # Step 1: Prepare the base image with status updates via the progress object
         base_image = prepare_base_image(uploaded_image, progress)
         progress(0.5, desc="Encrypting and embedding data...")
-        
-        # Step 2: Create the final encrypted image with overlays
-        created_image = create_encrypted_image(secret_data, public_key, base_image, show_keys)
-        
+        created_image = create_encrypted_image(secret_data, public_key, base_image, overlay_option)
         return created_image, f"✅ Success! Image created for '{service_name}'."
     except Exception as e:
         logger.error(f"Error creating image: {e}", exc_info=True)
@@ -230,21 +212,17 @@ def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoi
     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}'.")
-    
     server_url = endpoint_details['link']
     status = f"Connecting to remote server: {server_url}"
     yield None, status
-    
     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="/keylock-auth-decoder")
         decrypted_data = json.loads(result) if isinstance(result, str) else 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}"
@@ -285,7 +263,7 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
                     
                     gr.Markdown("### Image Options")
                     creator_base_image_input = gr.Image(label="Optional Base Image (600x600 recommended)", type="pil", sources=["upload"], show_download_button=False)
-                    creator_show_keys_checkbox = gr.Checkbox(label="Show key/value data on final image", value=True)
+                    creator_overlay_option = gr.Radio(label="Overlay Content", choices=["Keys and Values", "Keys Only", "None"], value="Keys and Values", info="Choose what data to display on the final image.")
                     
                     creator_button = gr.Button("✨ Create Auth Image", variant="primary", scale=2)
 
@@ -333,7 +311,7 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
     
     creator_button.click(
         fn=create_keylock_wrapper,
-        inputs=[creator_service_dropdown, creator_secret_input, endpoints_state, creator_base_image_input, creator_show_keys_checkbox],
+        inputs=[creator_service_dropdown, creator_secret_input, endpoints_state, creator_base_image_input, creator_overlay_option],
         outputs=[creator_image_output, creator_status]
     )