Update app.py

app.py

@@ -1,134 +1,27 @@
 import gradio as gr
-from gradio_client import Client
-from PIL import Image
-import base64
-import io
-import json
-import logging
-import requests
-import tempfile
-from cryptography.hazmat.primitives import serialization
-from cryptography.hazmat.primitives.asymmetric import rsa
-
-# --- Configure Logging ---
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# ==============================================================================
-#  CONFIGURATION: IDs AND URLs OF THE REMOTE SERVICES
-# ==============================================================================
-CREATOR_SPACE_ID = "broadfield-dev/KeyLock-Auth-Creator"
-SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server"
-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_URL = f"{BASE_HF_URL}{CREATOR_SPACE_ID}"
-SERVER_URL = f"{BASE_HF_URL}{SERVER_SPACE_ID}"
-CREATOR_APP_PY_URL = f"{CREATOR_URL}/blob/main/app.py"
-SERVER_APP_PY_URL = f"{SERVER_URL}/blob/main/app.py"
+import core # Import all our local logic
 
 # ==============================================================================
-#  API CALL WRAPPER FUNCTIONS (CORRECTED)
+#  GRADIO UI HELPER FUNCTIONS
 # ==============================================================================
-
-def get_creator_endpoints():
-    """Fetches the list of supported endpoints by making an HTTP request to the Creator's JSON file."""
-    status = f"Fetching endpoint list from {CREATOR_ENDPOINTS_JSON_URL}..."
-    # Using yield for streaming status updates
-    yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, [] 
+def creator_wrapper(public_key: str, secret_data: str):
+    """UI wrapper for the image creation logic."""
     try:
-        response = requests.get(CREATOR_ENDPOINTS_JSON_URL, timeout=10)
-        response.raise_for_status()
-        endpoints = response.json()
-        endpoint_names = [e['name'] for e in endpoints]
-        status = f"✅ Success! Found {len(endpoint_names)} endpoints."
-        yield gr.Dropdown(choices=endpoint_names, value=endpoint_names[0] if endpoint_names else None, label="Target Service"), status, endpoints
+        created_image = core.create_encrypted_image(secret_data, public_key)
+        return created_image, f"✅ Success! Image created. You can now go to the 'Send KeyLock' tab to test it."
     except Exception as e:
-        logger.error(f"Failed to get endpoints from creator's JSON file: {e}", exc_info=True)
-        status = f"❌ Error: Could not fetch configuration. Check the URL and if the 'endpoints.json' file is public. Details: {e}"
-        yield gr.Dropdown(choices=[], value=None, label="Error fetching services"), status, []
-
-def create_image_via_api(service_name: str, secret_data: str, available_endpoints: list):
-    """Calls the Creator Space API to generate an encrypted image for a selected service."""
-    if not all([service_name, secret_data]):
-        raise gr.Error("Please select a service and provide secret data.")
-    
-    status = f"Looking up public key for '{service_name}'..."
-    # No yield here, we'll return everything at the end.
+        return None, f"❌ Error: {e}"
 
+def decoder_wrapper(image: gr.Image, private_key: str):
+    """UI wrapper for the image decryption logic."""
     try:
-        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}' in the fetched configuration.")
-            
-        status = f"Connecting to Creator: {CREATOR_SPACE_ID}..."
-        logger.info(status)
-        
-        client = Client(src=CREATOR_SPACE_ID)
-        temp_filepath = client.predict(secret_data, public_key, api_name="/create_image")
-        
-        if not temp_filepath: 
-            raise gr.Error("Creator API did not return an image.")
-        
-        # --- PNG FIX ---
-        # Load the image from the temp path and return the PIL Image object directly.
-        # The gr.Image component's `format="png"` will handle the rest.
-        created_image = Image.open(temp_filepath)
-        
-        status = f"✅ Success! Image created for '{service_name}'."
-        return created_image, status
-
+        decrypted_data = core.decode_data_from_image(image, private_key)
+        return decrypted_data, "✅ Success! Data decrypted from image."
     except Exception as e:
-        logger.error(f"Creator API call failed: {e}", exc_info=True)
-        # On error, return an empty image and the error status.
-        return None, f"❌ Error calling Creator API: {e}"
-
-def decrypt_image_via_api(image: Image.Image):
-    """Calls the Server Space API to decrypt an image."""
-    if image is None: 
-        raise gr.Error("Please upload an image to decrypt.")
-        
-    status = f"Connecting to Server: {SERVER_SPACE_ID}..."
-    logger.info(status)
-    
-    try:
-        client = Client(src=SERVER_SPACE_ID)
-        with io.BytesIO() as buffer:
-            # Explicitly save as PNG to ensure format is correct before base64 encoding
-            image.save(buffer, format="PNG")
-            b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
-        
-        status = f"Calling API on {SERVER_SPACE_ID}..."
-        logger.info(status)
-        
-        decrypted_json = client.predict(b64_string, api_name="/keylock-auth-decoder")
-        
-        status = "✅ Success! Data decrypted by the Server."
-        logger.info(f"Decryption successful. Data: {decrypted_json}")
-        
-        # --- DECRYPTED DATA FIX ---
-        # Use a single, final return statement. This is more robust.
-        return decrypted_json, status
-
-    except Exception as e:
-        logger.error(f"Server API call failed: {e}", exc_info=True)
-        # Return an empty dict and the error message for a clean UI update.
-        return {}, f"❌ Error calling Server API: {e}"
-
-def generate_rsa_keys():
-    """Generates a new RSA 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
+        return None, f"❌ Error: {e}"
 
 # ==============================================================================
-#  GRADIO DASHBOARD INTERFACE (Corrected Layout and Components)
+#  GRADIO DASHBOARD INTERFACE
 # ==============================================================================
 theme = gr.themes.Base(
     primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky, neutral_hue=gr.themes.colors.slate,
@@ -140,77 +33,70 @@ theme = gr.themes.Base(
 )
 
 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, powered by live API calls to dedicated services.")
+    gr.Markdown("A self-contained demonstration of the entire KeyLock workflow: Key Generation, Image Creation, and Decryption.")
 
     with gr.Tabs() as tabs:
-        with gr.TabItem("① Auth Creator", id=0):
-            gr.Markdown("## Create an Encrypted Authentication Image")
-            gr.Markdown(f"This tool calls the **[{CREATOR_SPACE_ID}]({CREATOR_URL})** service to encrypt data for a chosen target. The list of targets is fetched live from the Creator's configuration.")
+        with gr.TabItem("① Generate Keys", id=0):
+            gr.Markdown("## Step 1: Create a Secure Key Pair")
+            gr.Markdown(
+                """
+                Every secure system starts with a key pair. This consists of a **Public Key** and a **Private Key**.
+                - **Public Key 🔑:** You can share this with anyone. It's used only for *encrypting* data. In a real system, you would add this key to a server's configuration file (like an `endpoints.json`) so that clients know how to encrypt data for it.
+                - **Private Key 🔐:** This must be kept **absolutely secret**. It is the only key that can *decrypt* data encrypted with its corresponding public key. In a real system, this would be stored as a secure environment variable or secret on the server.
+                """
+            )
             with gr.Row(variant="panel"):
+                with gr.Column(scale=1):
+                    gr.Markdown("### Your New Keys")
+                    gen_keys_button = gr.Button("Generate New 2048-bit Key Pair", variant="secondary")
+                    gr.Markdown("⬇️ **Copy these keys for the next steps!**")
                 with gr.Column(scale=2):
-                    with gr.Row():
-                        creator_service_dropdown = gr.Dropdown(label="Target Service", interactive=True, info="Select the API server you want 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 via API", variant="primary")
+                     with gr.Row():
+                        output_public_key = gr.Textbox(lines=11, label="Generated Public Key (For Creator)", interactive=False, show_copy_button=True)
+                        output_private_key = gr.Textbox(lines=11, label="Generated Private Key (For Decoder)", interactive=False, show_copy_button=True)
+
+        with gr.TabItem("② Create KeyLock", id=1):
+            gr.Markdown("## Step 2: Create an Encrypted Auth Image")
+            gr.Markdown("This tool acts as the **Auth Creator**. It takes your secret data and uses the **Public Key** from Step 1 to encrypt it into a new PNG image. This simulates a user preparing their credentials to send to a secure service.")
+            with gr.Row(variant="panel"):
+                with gr.Column(scale=1):
+                    gr.Markdown("### Configuration")
+                    creator_pubkey_input = gr.Textbox(lines=8, label="Paste the Public Key Here", placeholder="Copy the public key generated in Step 1...")
+                    creator_secret_input = gr.Textbox(lines=5, 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):
+                    gr.Markdown("### Output")
                     creator_status = gr.Textbox(label="Status", interactive=False, lines=2)
-                    # --- PNG FIX ---
-                    # The format='png' argument tells the component's download button to create a PNG.
-                    creator_image_output = gr.Image(label="Image from Creator Service", type="pil", show_download_button=True, format="png")
+                    # format="png" is the key to ensuring the download button creates a proper PNG file.
+                    creator_image_output = gr.Image(label="Generated Encrypted Image", type="pil", show_download_button=True, format="png", show_share_button=False)
 
-        with gr.TabItem("② Client / Decoder", id=1):
-            gr.Markdown("## Decrypt an Authentication Image")
-            gr.Markdown(f"This tool acts as a client, calling the **[{SERVER_SPACE_ID}]({SERVER_URL})** service to decrypt an image using its securely stored private key.")
+        with gr.TabItem("③ Send KeyLock", id=2):
+            gr.Markdown("## Step 3: Decrypt the Image (Client Simulation)")
+            gr.Markdown("This tool acts as the **Client** sending the image to a secure **Server**. To prove it can decrypt the data, the 'Server' needs the corresponding **Private Key** from Step 1. In a real application, you would only send the image; the server would already have its private key.")
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1):
-                    client_image_input = gr.Image(type="pil", label="Upload Encrypted Auth Image", sources=["upload", "clipboard"])
-                    client_button = gr.Button("🔓 Decrypt Image via Server API", variant="primary")
+                    gr.Markdown("### Input")
+                    client_image_input = gr.Image(type="pil", label="Upload or Drag Encrypted Image Here", sources=["upload", "clipboard"])
+                    client_private_key_input = gr.Textbox(lines=8, label="Paste the Private Key Here", placeholder="Copy the private key generated in Step 1...")
+                    client_button = gr.Button("🔓 Decrypt Image", variant="primary")
                 with gr.Column(scale=1):
+                    gr.Markdown("### Decrypted Data")
                     client_status = gr.Textbox(label="Status", interactive=False, lines=2)
-                    client_json_output = gr.JSON(label="Decrypted Data from Server")
-        
-        with gr.TabItem("ℹ️ Service Information", id=2):
-            gr.Markdown("## Ecosystem Architecture")
-            gr.Markdown("This dashboard coordinates separate Hugging Face Spaces to demonstrate a secure, decoupled workflow. Each service has a specific role.")
-            with gr.Row():
-                with gr.Column():
-                    gr.Markdown(f"### 🏭 Auth Creator Service\n- **Space:** [{CREATOR_SPACE_ID}]({CREATOR_URL})\n- **Role:** Provides an API to encrypt data for various targets defined in its `endpoints.json` file.\n- **Source Code:** [app.py]({CREATOR_APP_PY_URL})")
-                with gr.Column():
-                    gr.Markdown(f"### 📡 Decoder Server\n- **Space:** [{SERVER_SPACE_ID}]({SERVER_URL})\n- **Role:** The trusted authority. It holds a secret private key and provides a secure API to decrypt images.\n- **Source Code:** [app.py]({SERVER_APP_PY_URL})")
-            
-            # --- Generate Keys moved to a closed Accordion ---
-            with gr.Accordion("🔑 RSA Key Pair Generator", open=False):
-                gr.Markdown("Create a new public/private key pair. The public key can be added to a service's configuration to allow it to be a target for the Auth Creator.")
-                with gr.Row():
-                    with gr.Column():
-                        output_public_key = gr.Textbox(lines=10, label="Generated Public Key (Share This)", interactive=False, show_copy_button=True)
-                    with gr.Column():
-                        output_private_key = gr.Textbox(lines=10, label="Generated Private Key (Keep Secret!)", interactive=False, show_copy_button=True)
-                gen_keys_button = gr.Button("Generate New 2048-bit Key Pair", variant="secondary")
+                    client_json_output = gr.JSON(label="Result from 'Server'")
 
     # --- Wire up the component logic ---
-    gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
+    gen_keys_button.click(fn=core.generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
     
-    def refresh_endpoints():
-        # This is a generator function, so we need to iterate to get the last value.
-        *_, last_yield = get_creator_endpoints()
-        return last_yield
-
-    refresh_button.click(fn=refresh_endpoints, outputs=[creator_service_dropdown, creator_status, endpoints_state])
-    demo.load(fn=refresh_endpoints, outputs=[creator_service_dropdown, creator_status, endpoints_state])
-
     creator_button.click(
-        fn=create_image_via_api,
-        inputs=[creator_service_dropdown, creator_secret_input, endpoints_state],
-        outputs=[creator_image_output, creator_status] # Removed the gr.File output
+        fn=creator_wrapper,
+        inputs=[creator_pubkey_input, creator_secret_input],
+        outputs=[creator_image_output, creator_status]
     )
+    
     client_button.click(
-        fn=decrypt_image_via_api,
-        inputs=[client_image_input],
+        fn=decoder_wrapper,
+        inputs=[client_image_input, client_private_key_input],
         outputs=[client_json_output, client_status]
     )