Update app.py

app.py

@@ -1,220 +1,182 @@
 import gradio as gr
-import json
-import os
-import io
+from gradio_client import Client, GradioClientError
+from PIL import Image
 import base64
-import struct
+import io
 import logging
-import tempfile
-from PIL import Image, ImageDraw, ImageFont
-import numpy as np
-from cryptography.hazmat.primitives.ciphers.aead import AESGCM
-from cryptography.hazmat.primitives import hashes
-from cryptography.hazmat.primitives import serialization
-from cryptography.hazmat.primitives.asymmetric import padding
-from cryptography.exceptions import InvalidTag
-
-
-#import subprocess
-#result = subprocess.run(["pip", "freeze"], capture_output=True, text=True)
-#print(result.stdout)
-
 
 # --- Configure Logging ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 # ==============================================================================
-#  LOAD KEYS AND CONFIG
+#  CONFIGURATION: IDs of the remote Gradio Spaces
 # ==============================================================================
-KEYLOCK_PRIV_KEY = os.environ.get('KEYLOCK_PRIV_KEY')
-PUBLIC_KEY_PEM_STRING = ""
-try:
-    with open("keylock_pub.pem", "r") as f:
-        PUBLIC_KEY_PEM_STRING = f.read()
-    logger.info("Successfully loaded public key from keylock_pub.pem.")
-except FileNotFoundError:
-    PUBLIC_KEY_PEM_STRING = "Error: keylock_pub.pem file not found. 'Creator' tab will not work."
+# These should be the names of YOUR deployed spaces.
+CREATOR_SPACE_ID = "broadfield-dev/KeyLock-Auth-Creator"
+SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server"
+# We also provide a link to the original client for reference.
+CLIENT_SPACE_LINK = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-Client"
+
 
 # ==============================================================================
-#  CREATOR LOGIC
+#  API CALL WRAPPER FUNCTIONS
 # ==============================================================================
-def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.Image:
-    if not secret_data_str.strip(): raise ValueError("Secret data cannot be empty.")
-    
-    # 1. Parse & Encrypt Data
-    data_dict = {}
-    for line in secret_data_str.splitlines():
-        if not line.strip() or line.strip().startswith('#'): continue
-        key, value = line.split(':', 1) if ':' in line else line.split('=', 1)
-        data_dict[key.strip()] = value.strip().strip("'\"")
-
-    json_bytes = json.dumps(data_dict).encode('utf-8')
-    public_key = serialization.load_pem_public_key(public_key_pem.encode('utf-8'))
-    aes_key = os.urandom(32)
-    nonce = os.urandom(12)
-    ciphertext_with_tag = AESGCM(aes_key).encrypt(nonce, json_bytes, None)
-    rsa_encrypted_aes_key = public_key.encrypt(
-        aes_key, padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
-    )
-    encrypted_aes_key_len_bytes = struct.pack('>I', len(rsa_encrypted_aes_key))
-    encrypted_payload = encrypted_aes_key_len_bytes + rsa_encrypted_aes_key + nonce + ciphertext_with_tag
 
-    # 2. Create Carrier Image and Embed Data
-    img = Image.new('RGB', (800, 600), color=(45, 52, 54))
-    draw = ImageDraw.Draw(img)
-    try:
-        font = ImageFont.truetype("DejaVuSans.ttf", 40)
-    except IOError:
-        font = ImageFont.load_default()
-    draw.text((400, 300), "KeyLock Secure Data", fill=(223, 230, 233), font=font, anchor="ms")
+def create_image_via_api(secret_data: str, public_key: str):
+    """
+    Calls the Creator Space API to generate an encrypted image.
+    This function uses 'yield' to provide real-time status updates to the UI.
+    """
+    if not all([secret_data, public_key]):
+        raise gr.Error("Secret Data and Public Key are both required.")
     
-    img_rgb = img.convert("RGB")
-    pixel_data = np.array(img_rgb).ravel()
-    data_length_header = struct.pack('>I', len(encrypted_payload))
-    binary_payload = ''.join(format(byte, '08b') for byte in data_length_header + encrypted_payload)
-    if len(binary_payload) > pixel_data.size:
-        raise ValueError("Data is too large for the image capacity.")
-    for i in range(len(binary_payload)):
-        pixel_data[i] = (pixel_data[i] & 0xFE) | int(binary_payload[i])
-    stego_pixels = pixel_data.reshape(img_rgb.size[1], img_rgb.size[0], 3)
-    return Image.fromarray(stego_pixels, 'RGB')
+    status = f"Initializing client for Creator: {CREATOR_SPACE_ID}..."
+    yield None, None, status  # Yield initial status update
 
-# ==============================================================================
-#  SERVER (DECODER) LOGIC
-# ==============================================================================
-def decode_data_from_image(image: Image.Image) -> dict:
-    if not KEYLOCK_PRIV_KEY:
-        raise gr.Error("Server Error: The API is not configured with a private key.")
-    if image is None:
-        raise gr.Error("Please provide an image to decode.")
-    
     try:
-        private_key = serialization.load_pem_private_key(KEYLOCK_PRIV_KEY.encode(), password=None)
-        with io.BytesIO() as buffer:
-            image.save(buffer, format="PNG")
-            image_buffer = buffer.getvalue()
-        
-        img_rgb = Image.open(io.BytesIO(image_buffer)).convert("RGB")
-        pixel_data = np.array(img_rgb).ravel()
+        client = Client(src=CREATOR_SPACE_ID)
+        status = f"Calling API '/create_image' on {CREATOR_SPACE_ID}..."
+        yield None, None, status
         
-        header_binary_string = "".join(str(p & 1) for p in pixel_data[:32])
-        data_length = int(header_binary_string, 2)
-        end_offset = 32 + (data_length * 8)
-        if pixel_data.size < end_offset: raise ValueError("Image data corrupt or truncated.")
+        # The 'predict' method calls the API endpoint.
+        # The result from an Image output in the remote API is a filepath to a temporary file.
+        temp_filepath = client.predict(
+            secret_data_str=secret_data,
+            public_key_pem=public_key,
+            api_name="/create_image"  # The API name defined in the Creator space
+        )
         
-        data_binary_string = "".join(str(p & 1) for p in pixel_data[32:end_offset])
-        crypto_payload = int(data_binary_string, 2).to_bytes(data_length, byteorder='big')
+        if not temp_filepath:
+            raise gr.Error("Creator API did not return a valid image file path.")
+            
+        # Load the image from the temporary file to display it in our dashboard's UI
+        created_image = Image.open(temp_filepath)
         
-        rsa_key_len = private_key.key_size // 8
-        len_header = struct.unpack('>I', crypto_payload[:4])[0]
-        if len_header != rsa_key_len: raise ValueError("Key pair mismatch or corrupt data.")
+        status = "✅ Success! Image created by the Creator service."
+        yield created_image, temp_filepath, status
+
+    except Exception as e:
+        logger.error(f"Creator API call failed: {e}", exc_info=True)
+        # Yield the error message back to the UI
+        yield None, None, f"❌ Error calling Creator API: {e}"
+
+
+def decrypt_image_via_api(image: Image.Image):
+    """
+    Calls the Server Space API to decrypt an image.
+    Uses 'yield' for status updates.
+    """
+    if image is None:
+        raise gr.Error("Please upload an image to decrypt.")
+
+    status = f"Initializing client for Server: {SERVER_SPACE_ID}..."
+    yield None, status
+
+    try:
+        client = Client(src=SERVER_SPACE_ID)
         
-        encrypted_key = crypto_payload[4:4 + rsa_key_len]
-        nonce = crypto_payload[4 + rsa_key_len : 4 + rsa_key_len + 12]
-        ciphertext = crypto_payload[4 + rsa_key_len + 12:]
+        # Convert the user's uploaded PIL image to a base64 string for the server's API
+        with io.BytesIO() as buffer:
+            image.save(buffer, format="PNG")
+            b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
         
-        aes_key = private_key.decrypt(encrypted_key, padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None))
-        decrypted_bytes = AESGCM(aes_key).decrypt(nonce, ciphertext, None)
+        status = f"Calling API '/keylock-auth-decoder' on {SERVER_SPACE_ID}..."
+        yield None, status
         
-        logger.info("Successfully decoded an image via internal call.")
-        return json.loads(decrypted_bytes.decode('utf-8'))
+        # Call the server's API endpoint. The result will be a dictionary.
+        decrypted_json = client.predict(
+            image_base64_string=b64_string,
+            api_name="/keylock-auth-decoder"  # The API name defined in the Server space
+        )
         
-    except Exception as e:
-        logger.error(f"Decryption failed: {e}", exc_info=True)
-        raise gr.Error(f"Decryption failed. The image may be corrupt, not encrypted with the correct key, or the server is misconfigured. Details: {e}")
+        status = "✅ Success! Data decrypted by the Server."
+        yield decrypted_json, status
 
-# ==============================================================================
-#  GRADIO UI HELPER FUNCTIONS
-# ==============================================================================
-def creator_wrapper(secret_data_str: str):
-    """Wrapper for the Gradio UI to call the creator function."""
-    if not PUBLIC_KEY_PEM_STRING or "Error" in PUBLIC_KEY_PEM_STRING:
-        raise gr.Error("Cannot create image: public key file is missing or invalid.")
-    try:
-        encrypted_image = create_encrypted_image(secret_data_str, PUBLIC_KEY_PEM_STRING)
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
-            encrypted_image.save(tmp.name)
-            return encrypted_image, tmp.name, "✅ Success! Image created and ready for download."
     except Exception as e:
-        return None, None, f"❌ Error: {e}"
+        logger.error(f"Server API call failed: {e}", exc_info=True)
+        yield None, f"❌ Error calling Server API: {e}"
+
 
 # ==============================================================================
 #  GRADIO DASHBOARD INTERFACE
 # ==============================================================================
 theme = gr.themes.Base(
     primary_hue="blue",
-    secondary_hue="blue",
+    secondary_hue="sky",
     neutral_hue="slate",
     font=(gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"),
 ).set(
-    body_background_fill="#F0F2F5",
+    body_background_fill="#F8FAFC",
     block_background_fill="white",
     block_border_width="1px",
     block_shadow="*shadow_drop_lg",
-    button_primary_background_fill="*primary_500",
-    button_primary_background_fill_hover="*primary_600",
+    button_primary_background_fill="*primary_600",
+    button_primary_background_fill_hover="*primary_700",
     button_primary_text_color="white",
 )
 
 with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
     gr.Markdown("# 🔑 KeyLock Operations Dashboard")
-    gr.Markdown("A unified dashboard to demonstrate the KeyLock ecosystem: API Server, Image Creator, and API Client.")
+    gr.Markdown("A centralized dashboard demonstrating the entire KeyLock ecosystem by making live API calls to the deployed Creator and Server spaces.")
 
     with gr.Tabs():
-        with gr.TabItem("📡 API Server Status & Docs", id=0):
-            gr.Markdown("## Server Status")
-            key_status = "✅ Loaded successfully from secrets." if KEYLOCK_PRIV_KEY else "❌ NOT FOUND. The API will not work. Set the `KEYLOCK_PRIV_KEY` secret in the Space settings."
-            gr.Textbox(label="Private Key Status", value=key_status, interactive=False)
-            
-            with gr.Accordion("View Service Public Key", open=False):
-                gr.Code(value=PUBLIC_KEY_PEM_STRING, language="python", label="Service Public Key (from keylock_pub.pem)")
-            
-            gr.Markdown("## API Documentation")
-            gr.Markdown("This server exposes a conceptual API endpoint for external clients.")
-            gr.Code(language="markdown", value="""
-- **Endpoint**: `/run/keylock-auth-decoder`
-- **Method**: `POST`
-- **Body**: JSON payload `{ "data": ["<base64_encoded_image_string>"] }`
-""")
-
-        with gr.TabItem("🏭 Image Creator", id=1):
-            gr.Markdown("## Create an Encrypted Image for the Server")
-            gr.Markdown("Use this tool to encrypt data using the server's public key.")
-            with gr.Row():
+        with gr.TabItem("🏭 Image Creator", id=0):
+            gr.Markdown("## Create an Encrypted Image via API Call")
+            gr.Markdown(f"This interface calls the **`{CREATOR_SPACE_ID}`** Space to generate a new encrypted image.")
+            with gr.Row(variant="panel"):
                 with gr.Column(scale=2):
-                    creator_input = gr.Textbox(
-                        lines=7,
-                        label="Secret Data (Key-Value Pairs)",
-                        placeholder="USERNAME: [email protected]\nAPI_KEY: sk-12345..."
-                    )
-                    creator_button = gr.Button("Create Encrypted Image", variant="primary")
+                    creator_secret_input = gr.Textbox(lines=5, label="Secret Data", placeholder="API_KEY: sk-123...\nUSER: demo-user")
+                    creator_pubkey_input = gr.Textbox(lines=7, label="Public Key of the Target Server", placeholder="Paste the Server's public key here...")
+                    creator_button = gr.Button("Create Image via Creator API", variant="primary", icon="✨")
                 with gr.Column(scale=1):
-                    creator_status = gr.Textbox(label="Status", interactive=False)
-                    creator_image_output = gr.Image(label="Generated Encrypted Image", type="pil")
-                    creator_download_output = gr.File(label="Download Image")
-
-        with gr.TabItem("💻 API Client (Simulator)", id=2):
-            gr.Markdown("## Decrypt an Image via Simulated API Call")
-            gr.Markdown("Upload an image created by the 'Creator' tab to test the server's decryption logic.")
-            with gr.Row():
+                    creator_status = gr.Textbox(label="Status", interactive=False, lines=2)
+                    creator_image_output = gr.Image(label="Image from Creator Service", type="pil", show_download_button=True)
+                    creator_download_output = gr.File(label="Download Image File")
+
+        with gr.TabItem("💻 Client / Decoder", id=1):
+            gr.Markdown("## Decrypt an Image via API Call")
+            gr.Markdown(f"This interface acts as a client, calling the **`{SERVER_SPACE_ID}`** Space to decrypt an image.")
+            with gr.Row(variant="panel"):
                 with gr.Column(scale=1):
-                    client_input_image = gr.Image(type="pil", label="Upload Encrypted Image", sources=["upload", "clipboard"])
-                    client_decrypt_button = gr.Button("Decrypt Image", variant="primary")
+                    client_image_input = gr.Image(type="pil", label="Upload Encrypted Image", sources=["upload", "clipboard"])
+                    client_button = gr.Button("Decrypt Image via Server API", variant="primary", icon="🔓")
                 with gr.Column(scale=1):
-                    client_output_json = gr.JSON(label="Decrypted Data from Server")
+                    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 Overview")
+            gr.Markdown(
+                f"""
+                This dashboard coordinates three separate Hugging Face Spaces to demonstrate a complete workflow:
+
+                1.  **Creator Service:**
+                    -   **Space:** [{CREATOR_SPACE_ID}](https://huggingface.co/spaces/{CREATOR_SPACE_ID})
+                    -   **Role:** Provides a UI and an API (`/run/create_image`) to encrypt data into PNG images.
+
+                2.  **Server (Decoder) Service:**
+                    -   **Space:** [{SERVER_SPACE_ID}](https://huggingface.co/spaces/{SERVER_SPACE_ID})
+                    -   **Role:** Holds a secret private key and provides a secure API (`/run/keylock-auth-decoder`) to decrypt images.
+
+                3.  **This Dashboard:**
+                    -   **Role:** Acts as a master client and control panel, making live API calls to the other services to showcase the end-to-end process.
+                
+                *Note: The original `{CLIENT_SPACE_LINK.split('/')[-1]}` is now superseded by the 'Client / Decoder' tab in this dashboard.*
+                """
+            )
 
     # --- Wire up the component logic ---
     creator_button.click(
-        fn=creator_wrapper,
-        inputs=[creator_input],
+        fn=create_image_via_api,
+        inputs=[creator_secret_input, creator_pubkey_input],
         outputs=[creator_image_output, creator_download_output, creator_status]
     )
     
-    client_decrypt_button.click(
-        fn=decode_data_from_image,
-        inputs=[client_input_image],
-        outputs=[client_output_json]
+    client_button.click(
+        fn=decrypt_image_via_api,
+        inputs=[client_image_input],
+        outputs=[client_json_output, client_status]
     )
 
 if __name__ == "__main__":