Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image, ImageDraw, ImageFont, ImageOps
 import base64
 import io
 import json
@@ -13,10 +13,12 @@ from cryptography.hazmat.primitives.asymmetric import rsa, padding
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from cryptography.hazmat.primitives import hashes
 from gradio_client import Client
+from huggingface_hub import InferenceClient
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+# --- Constants ---
 CREATOR_ENDPOINTS_JSON_URL = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-Creator/raw/main/endpoints.json"
 BASE_HF_URL = "https://huggingface.co/spaces/"
 CREATOR_SPACE_ID = "broadfield-dev/KeyLock-Auth-Creator"
@@ -26,21 +28,77 @@ SERVER_URL = f"{BASE_HF_URL}{SERVER_SPACE_ID}"
 CREATOR_APP_PY_URL = f"{CREATOR_URL}/raw/main/app.py"
 SERVER_APP_PY_URL = f"{SERVER_URL}/raw/main/app.py"
 
+# NEW: Constants for improved image generation
+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 ---
+
+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:
+    """
+    Provides a base image using the fallback logic, updating a Gradio progress object.
+    1. User-uploaded image.
+    2. Default URL image.
+    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}")
+        response = requests.get(DEFAULT_IMAGE_URL, timeout=15)
+        response.raise_for_status()
+        img = Image.open(io.BytesIO(response.content)).convert("RGB")
+        progress(0, desc="✅ Using default background image.")
+        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}")
+        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.")
+        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 ---
 
 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')
+    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:
+def create_encrypted_image(
+    secret_data_str: str,
+    public_key_pem: str,
+    base_image: Image.Image,
+    show_key_overlay: bool
+) -> Image.Image:
     if not secret_data_str.strip():
         raise ValueError("Secret data cannot be empty.")
     if not public_key_pem.strip():
@@ -49,8 +107,7 @@ def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.I
     data_dict = {}
     for line in secret_data_str.splitlines():
         line = line.strip()
-        if not line or line.startswith('#'):
-            continue
+        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("'\"")
@@ -64,32 +121,73 @@ def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.I
     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)
-    )
+    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)
+    # 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
     try:
-        font = ImageFont.truetype("DejaVuSans.ttf", 40)
+        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 = ImageFont.load_default(size=30)
-    draw.text((400, 300), "KeyLock Secure Data", fill=(223, 230, 233), font=font, anchor="ms")
-    
+        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
+
+    # 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:
+        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]
+        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 ---
     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.")
+        raise ValueError(f"Data is too large for the image. Max size: {pixel_data.size // 8} bytes. Your data: ~{len(binary_payload) // 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((600, 800, 3))
+    stego_pixels = pixel_data.reshape((height, width, 3))
     
     return Image.fromarray(stego_pixels, 'RGB')
 
+# --- Gradio Wrapper Functions ---
+
 def get_server_list():
     status = "Fetching server list from remote config..."
     yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, []
@@ -97,62 +195,43 @@ def get_server_list():
         response = requests.get(CREATOR_ENDPOINTS_JSON_URL, timeout=10)
         response.raise_for_status()
         all_entries = response.json()
-        valid_endpoints = []
-        
-        for entry in all_entries:
-            if not isinstance(entry, dict) or "name" not in entry or "public_key" not in entry:
-                logger.warning(f"Skipping invalid entry (missing name or public_key): {entry}")
-                continue
-
-            if "api_endpoint" not in entry:
-                if "link" in entry:
-                    base_url = entry["link"].strip("/")
-                    entry["api_endpoint"] = f"{base_url}/run/predict/"
-                    logger.info(f"Auto-generating API endpoint for '{entry['name']}': {entry['api_endpoint']}")
-                else:
-                    logger.warning(f"Skipping invalid entry '{entry['name']}' (missing 'api_endpoint' and 'link'): {entry}")
-                    continue
-
-            valid_endpoints.append(entry)
-
+        valid_endpoints = [e for e in all_entries if isinstance(e, dict) and "name" in e and "public_key" in e and ("api_endpoint" in e or "link" in e)]
         if not valid_endpoints:
-            raise ValueError("No valid server configurations found in the remote file.")
-            
+            raise ValueError("No valid server configurations found.")
         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}"
+        status = f"❌ Error fetching 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):
+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)):
     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:
-        created_image = create_encrypted_image(secret_data, public_key)
+        # 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)
+        
         return created_image, f"✅ Success! Image created for '{service_name}'."
     except Exception as e:
-        logger.error(f"Error creating image: {e}")
+        logger.error(f"Error creating image: {e}", exc_info=True)
         return None, f"❌ Error: {e}"
 
 def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoints: 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.")
-        
+    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"Configuration Error: Could not find details for '{service_name}'. Refresh the list.")
+    if not endpoint_details or not endpoint_details.get('link'): raise gr.Error(f"Config Error for '{service_name}'.")
     
-    server_url = endpoint_details.get('link')
-    if not server_url:
-        raise gr.Error(f"Configuration Error: The selected server '{service_name}' is missing a 'link' to its Space.")
-
+    server_url = endpoint_details['link']
     status = f"Connecting to remote server: {server_url}"
     yield None, status
     
@@ -162,28 +241,20 @@ def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoi
             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 = result
-        if isinstance(decrypted_data, str):
-            try:
-                decrypted_data = json.loads(decrypted_data)
-            except json.JSONDecodeError:
-                pass
+        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}")
+        logger.error(f"Error calling server with gradio_client: {e}", exc_info=True)
         yield None, f"❌ Error calling server API: {e}"
 
 def refresh_and_update_all():
-    for dropdown_update, status_update, state_update in get_server_list():
-        pass
+    for dropdown_update, status_update, state_update in get_server_list(): pass
     return dropdown_update, dropdown_update, status_update, state_update
 
+# --- Gradio UI ---
+
 theme = gr.themes.Base(
     primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky, neutral_hue=gr.themes.colors.slate,
     font=(gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"),
@@ -202,22 +273,28 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
     with gr.Tabs() as tabs:
         with gr.TabItem("① Create KeyLock", id=0):
             gr.Markdown("## Step 1: Create an Encrypted Authentication Image (Local)")
-            gr.Markdown("# ❗ DEMO: Don't send personal information encrypted to the demo server")
-            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"Encrypt your data into a PNG. You can upload your own background image, or we'll provide a stylish default or generate one for you. The encryption process happens entirely in your browser.")
             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")
-                    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):
-                    creator_status = gr.Textbox(label="Status", interactive=False, lines=2)
-                    creator_image_output = gr.Image(label="Generated Encrypted Image", type="pil", show_download_button=True, format="png")
+                    
+                    creator_secret_input = gr.Textbox(lines=5, label="Secret Data to Encrypt", placeholder="API_KEY: sk-123...\nUSER: demo-user")
+                    
+                    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_button = gr.Button("✨ Create Auth Image", variant="primary", scale=2)
+
+                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, format="png", height=600)
 
         with gr.TabItem("② Send KeyLock", id=1):
             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.")
+            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. The server uses its securely stored private key to decrypt the data and sends the result back.")
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1):
                     gr.Markdown("### Configuration")
@@ -248,10 +325,17 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
                         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", variant="secondary")
 
+    # --- Event Handlers ---
     gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
     refresh_button.click(fn=refresh_and_update_all, outputs=[creator_service_dropdown, send_service_dropdown, creator_status, endpoints_state])
     demo.load(fn=refresh_and_update_all, outputs=[creator_service_dropdown, send_service_dropdown, creator_status, endpoints_state])
-    creator_button.click(fn=create_keylock_wrapper, inputs=[creator_service_dropdown, creator_secret_input, endpoints_state], outputs=[creator_image_output, creator_status])
+    
+    creator_button.click(
+        fn=create_keylock_wrapper,
+        inputs=[creator_service_dropdown, creator_secret_input, endpoints_state, creator_base_image_input, creator_show_keys_checkbox],
+        outputs=[creator_image_output, creator_status]
+    )
+    
     client_button.click(fn=send_keylock_wrapper, inputs=[send_service_dropdown, client_image_input, endpoints_state], outputs=[client_json_output, client_status])
 
 if __name__ == "__main__":