Create client/app.py

client/app.py

@@ -0,0 +1,320 @@
+# 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
+import io
+import json
+import logging
+import os
+import requests
+import struct
+import numpy as np
+from cryptography.hazmat.primitives import serialization
+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, handle_file
+from huggingface_hub import InferenceClient
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+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) ---
+# ...
+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:
+        return ImageOps.fit(img, (size, size), Image.Resampling.LANCZOS)
+    except Exception as e:
+        logger.error(f"Failed to resize and crop image: {e}")
+        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.
+    """
+    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}")
+        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.")
+    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 (Your original code is perfect here) ---
+# ...
+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')
+    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:
+    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.")
+    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
+
+    # Add overlays
+    draw = ImageDraw.Draw(img, "RGBA")
+    try:
+        font_bold = ImageFont.truetype("DejaVuSans-Bold.ttf", 30)
+        font_regular = ImageFont.truetype("DejaVuSans.ttf", 15)
+    except IOError:
+        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)
+
+    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()]
+
+        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
+
+        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()
+    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.")
+    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.")
+        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 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.")
+    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:
+        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:
+        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):
+    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}'.")
+
+    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
+
+    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
+        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
+
+        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:
+    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.")
+
+    with gr.Tabs() as tabs:
+        with gr.TabItem("① Create KeyLock", id=0):
+            # ... UI components for creating
+            with gr.Row():
+                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_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):
+             # ... UI components for sending
+            with gr.Row():
+                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_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")
+
+        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")
+
+    # --- Event Handlers ---
+    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
+
+    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
+    )
+
+    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],
+        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)
+
+if __name__ == "__main__":
+    demo.launch()