Update app.py

app.py

@@ -1,24 +1,130 @@
 import gradio as gr
-import core # Import all our local logic
+import json
+import os
+import io
+import base64
+import struct
+import logging
+import requests
+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 rsa, padding
+from cryptography.exceptions import InvalidTag
+
+# --- Configure Logging ---
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 # ==============================================================================
-#  GRADIO UI HELPER FUNCTIONS
+#  CONFIGURATION: URL of the Remote SERVER Service
+# ==============================================================================
+SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server"
+BASE_HF_URL = "https://huggingface.co/spaces/"
+SERVER_URL = f"{BASE_HF_URL}{SERVER_SPACE_ID}"
+# The API endpoint is constructed from the server's direct URL, not the hub URL.
+SERVER_DIRECT_URL_BASE = f"https://{SERVER_SPACE_ID.replace('/', '-')}.hf.space"
+SERVER_API_ENDPOINT = f"{SERVER_DIRECT_URL_BASE}/run/keylock-auth-decoder"
+
+# ==============================================================================
+#  LOCAL LOGIC (Key and Image Generation)
 # ==============================================================================
-def creator_wrapper(public_key: str, secret_data: str):
-    """UI wrapper for the image creation logic."""
-    try:
-        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:
-        return None, f"❌ Error: {e}"
 
-def decoder_wrapper(image: gr.Image, private_key: str):
-    """UI wrapper for the image decryption logic."""
+def generate_rsa_keys():
+    """Generates a new 2048-bit RSA key pair LOCALLY."""
+    logger.info("Generating new RSA key pair locally.")
+    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) -> Image.Image:
+    """Creates the encrypted image LOCALLY."""
+    logger.info("Starting local image creation process...")
+    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():
+        if not line.strip() or line.strip().startswith('#'): continue
+        parts = line.split(':', 1) if ':' in line else line.split('=', 1)
+        if len(parts) != 2: continue
+        data_dict[parts[0].strip()] = parts[1].strip().strip("'\"")
+    if not data_dict: raise ValueError("No valid key-value pairs found.")
+    
+    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_with_tag = AESGCM(aes_key).encrypt(nonce, json_bytes, None)
+    rsa_encrypted_aes_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_aes_key)) + rsa_encrypted_aes_key + nonce + ciphertext_with_tag
+    
+    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(size=30)
+    draw.text((400, 300), "KeyLock Secure Data", fill=(223, 230, 233), font=font, anchor="ms")
+    
+    pixel_data = np.array(img.convert("RGB")).ravel()
+    binary_payload = ''.join(format(byte, '08b') for byte in struct.pack('>I', len(encrypted_payload)) + encrypted_payload)
+    if len(binary_payload) > pixel_data.size: raise ValueError("Data too large for 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((600, 800, 3))
+    return Image.fromarray(stego_pixels, 'RGB')
+
+# ==============================================================================
+#  REMOTE API CALL LOGIC
+# ==============================================================================
+
+def decrypt_image_via_api(image: Image.Image):
+    """Makes a LIVE API call to the deployed server to decrypt an image."""
+    if image is None: raise gr.Error("Please provide an image to send.")
+        
+    status = f"Connecting to server: {SERVER_SPACE_ID}..."
+    yield None, status
+    
     try:
-        decrypted_data = core.decode_data_from_image(image, private_key)
-        return decrypted_data, "✅ Success! Data decrypted from image."
+        with io.BytesIO() as buffer:
+            image.save(buffer, format="PNG")
+            b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        
+        payload = {"data": [b64_string]}
+        headers = {"Content-Type": "application/json"}
+        
+        status = f"Sending image to API endpoint:\n{SERVER_API_ENDPOINT}"
+        yield None, status
+        
+        response = requests.post(SERVER_API_ENDPOINT, headers=headers, json=payload, timeout=45)
+        response_json = response.json()
+        
+        if response.status_code == 200:
+            if "data" in response_json:
+                decrypted_data = response_json["data"][0]
+                status = "✅ Success! Data decrypted by the remote server."
+                return decrypted_data, status
+            elif "error" in response_json:
+                raise gr.Error(f"API returned an error: {response_json['error']}")
+        else:
+            error_detail = response_json.get("error", "Unknown error.")
+            raise gr.Error(f"API Error (Status {response.status_code}): {error_detail}")
+
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Network error calling API: {e}")
+        raise gr.Error(f"Could not connect to the API. Check the server space is running and the URL is correct. Error: {e}")
     except Exception as e:
-        return None, f"❌ Error: {e}"
+        logger.error(f"An unexpected error occurred: {e}", exc_info=True)
+        raise gr.Error(f"An unexpected error occurred: {e}")
+
 
 # ==============================================================================
 #  GRADIO DASHBOARD INTERFACE
@@ -34,69 +140,80 @@ theme = gr.themes.Base(
 
 with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
     gr.Markdown("# 🔑 KeyLock Operations Dashboard")
-    gr.Markdown("A self-contained demonstration of the entire KeyLock workflow: Key Generation, Image Creation, and Decryption.")
+    gr.Markdown("A self-contained dashboard to demonstrate the 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("① Generate Keys", id=0):
-            gr.Markdown("## Step 1: Create a Secure Key Pair")
+            gr.Markdown("## Step 1: Create a Secure Key Pair (Local)")
             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.
+                This tool generates a new RSA key pair within your browser session. In a real-world scenario, the **Private Key** would be immediately stored as a secure secret on a server (like the `KEYLOCK_PRIV_KEY` secret on our demo server), and would never be shown in a UI like this. The **Public Key** would be distributed to clients or other services that need to encrypt data for that server.
+                
+                **Action:** Click the button below, then copy both keys for the next steps.
                 """
             )
             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!**")
+                    gen_keys_button = gr.Button("Generate New 2048-bit Key Pair", icon="🔑", variant="secondary")
                 with gr.Column(scale=2):
                      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.")
+            gr.Markdown("## Step 2: Create an Encrypted Auth Image (Local)")
+            gr.Markdown(
+                """
+                This tool acts as the **Auth Creator**. It takes your secret data and uses the **Public Key** you generated in Step 1 to encrypt it into a new PNG image. This entire process happens locally in this application. This simulates a user or an automated client preparing credentials to send to the secure server.
+
+                **Action:** Paste the **Public Key** from Step 1, enter some secrets, and click create.
+                """
+            )
             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_secret_input = gr.Textbox(lines=5, label="Secret Data to Encrypt", placeholder="SESSION_ID: abc-123\nUSER: demo@example.com")
                     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)
-                    # 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("③ 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.")
+            gr.Markdown("## Step 3: Decrypt via Live API Call")
+            gr.Markdown(
+                f"""
+                This is the core demonstration. This tool acts as a **Client** sending the encrypted image to our live, remote **Server** at [{SERVER_SPACE_ID}]({SERVER_URL}). 
+                
+                For this demo to work, the **Private Key** you generated in Step 1 must be **the same one** set as the `KEYLOCK_PRIV_KEY` secret in the `{SERVER_SPACE_ID}` Space settings. The client sends the image, and the server uses its own secret key to decrypt it.
+
+                **Action:** Upload the image from Step 2. The dashboard will make a live API call to `{SERVER_API_ENDPOINT}`.
+                """
+            )
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1):
                     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")
+                    client_button = gr.Button("🔓 Decrypt Image via Remote Server", 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="Result from 'Server'")
+                    client_json_output = gr.JSON(label="Result from Server")
 
     # --- Wire up the component logic ---
-    gen_keys_button.click(fn=core.generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
+    gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
     
     creator_button.click(
-        fn=creator_wrapper,
+        fn=create_encrypted_image,
         inputs=[creator_pubkey_input, creator_secret_input],
         outputs=[creator_image_output, creator_status]
     )
     
     client_button.click(
-        fn=decoder_wrapper,
-        inputs=[client_image_input, client_private_key_input],
+        fn=decrypt_image_via_api,
+        inputs=[client_image_input],
         outputs=[client_json_output, client_status]
     )