Update app.py

app.py

@@ -12,12 +12,11 @@ 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
 
-# --- Basic Configuration ---
 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"
@@ -29,7 +28,6 @@ SERVER_APP_PY_URL = f"{SERVER_URL}/raw/main/app.py"
 
 
 def generate_rsa_keys():
-    """Generates a new 2048-bit RSA private and public key pair."""
     private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
     private_pem = private_key.private_bytes(
         encoding=serialization.Encoding.PEM,
@@ -43,7 +41,6 @@ def generate_rsa_keys():
     return private_pem, public_pem
 
 def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.Image:
-    """Encrypts data and embeds it into a new image using steganography."""
     if not secret_data_str.strip():
         raise ValueError("Secret data cannot be empty.")
     if not public_key_pem.strip():
@@ -64,20 +61,16 @@ def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.I
     json_bytes = json.dumps(data_dict).encode('utf-8')
     public_key = serialization.load_pem_public_key(public_key_pem.encode('utf-8'))
     
-    # AES-GCM encryption
     aes_key, nonce = os.urandom(32), os.urandom(12)
     ciphertext = AESGCM(aes_key).encrypt(nonce, json_bytes, None)
     
-    # RSA-OAEP encryption for the AES key
     rsa_encrypted_key = public_key.encrypt(
         aes_key,
         padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
     )
     
-    # Pack the payload: [len(rsa_key)][rsa_key][nonce][ciphertext]
     encrypted_payload = struct.pack('>I', len(rsa_encrypted_key)) + rsa_encrypted_key + nonce + ciphertext
     
-    # Create a base image
     img = Image.new('RGB', (800, 600), color=(45, 52, 54))
     draw = ImageDraw.Draw(img)
     try:
@@ -86,7 +79,6 @@ def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.I
         font = ImageFont.load_default(size=30)
     draw.text((400, 300), "KeyLock Secure Data", fill=(223, 230, 233), font=font, anchor="ms")
     
-    # LSB Steganography
     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)
     
@@ -99,7 +91,6 @@ def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.I
     return Image.fromarray(stego_pixels, 'RGB')
 
 def get_server_list():
-    """Fetches and validates the server list from the remote JSON configuration."""
     status = "Fetching server list from remote config..."
     yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, []
     try:
@@ -113,7 +104,6 @@ def get_server_list():
                 logger.warning(f"Skipping invalid entry (missing name or public_key): {entry}")
                 continue
 
-            # **FIX: Auto-generate api_endpoint from link if it's missing**
             if "api_endpoint" not in entry:
                 if "link" in entry:
                     base_url = entry["link"].strip("/")
@@ -137,7 +127,6 @@ def get_server_list():
         yield gr.Dropdown(choices=[], value=None, label="Error fetching servers"), status, []
 
 def create_keylock_wrapper(service_name: str, secret_data: str, available_endpoints: list):
-    """Wrapper function to handle UI for creating the encrypted image."""
     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)
@@ -151,17 +140,20 @@ def create_keylock_wrapper(service_name: str, secret_data: str, available_endpoi
         return None, f"❌ Error: {e}"
 
 def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoints: list):
-    """Wrapper function to handle UI for sending the image to the remote server."""
     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.")
         
-    api_endpoint = next((e.get('api_endpoint') for e in available_endpoints if e['name'] == service_name), None)
-    if not api_endpoint:
-        raise gr.Error(f"Configuration Error: Could not find 'api_endpoint' for '{service_name}'.")
-        
-    status = f"Connecting to endpoint: {api_endpoint}"
+    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.")
+    
+    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.")
+
+    status = f"Connecting to remote server: {server_url}"
     yield None, status
     
     try:
@@ -169,40 +161,29 @@ def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoi
             image.save(buffer, format="PNG")
             b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
         
-        # Gradio API expects data in this format
-        payload = {"data": [b64_string]}
-        headers = {"Content-Type": "application/json"}
-        
-        response = requests.post(api_endpoint, headers=headers, json=payload, timeout=45)
-        response.raise_for_status() # **IMPROVEMENT: Raise exception for HTTP errors**
+        client = Client(server_url)
+        result = client.predict(
+                image_base64_string=b64_string,
+                api_name="/keylock-auth-decoder"
+        )
         
-        response_json = response.json()
-        if "data" in response_json:
-            decrypted_data = response_json["data"][0]
-            # If the server returns a stringified JSON, parse it again
-            if isinstance(decrypted_data, str):
-                try:
-                    decrypted_data = json.loads(decrypted_data)
-                except json.JSONDecodeError:
-                    pass # Keep as string if not valid JSON
-            yield decrypted_data, "✅ Success! Data decrypted by remote server."
-        elif "error" in response_json:
-            raise gr.Error(f"API Error: {response_json['error']}")
-        else:
-            raise gr.Error(f"API returned an unexpected response format: {response_json}")
+        decrypted_data = result
+        if isinstance(decrypted_data, str):
+            try:
+                decrypted_data = json.loads(decrypted_data)
+            except json.JSONDecodeError:
+                pass
+        yield decrypted_data, "✅ Success! Data decrypted by remote server."
 
     except Exception as e:
-        logger.error(f"Error calling server API: {e}")
+        logger.error(f"Error calling server with gradio_client: {e}")
         yield None, f"❌ Error calling server API: {e}"
 
 def refresh_and_update_all():
-    """Calls get_server_list and updates all relevant UI components."""
-    # The generator yields updates, we need to iterate to get the final one
     for dropdown_update, status_update, state_update in get_server_list():
         pass
     return dropdown_update, dropdown_update, status_update, state_update
 
-# --- Gradio UI Definition ---
 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"),
@@ -235,7 +216,6 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
 
         with gr.TabItem("② Send KeyLock", id=1):
             gr.Markdown("## Step 2: Decrypt via Live API Call")
-            gr.Markdown("## This a Demo, don't send personal encrypted secret data to the demo server.")
             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.")
             with gr.Row(variant="panel"):
                 with gr.Column(scale=1):
@@ -267,7 +247,6 @@ 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])