Update app.py

app.py

@@ -29,13 +29,7 @@ except FileNotFoundError:
     logger.error("FATAL: keylock_pub.pem not found.")
     PUBLIC_KEY_PEM_STRING = "Error: keylock_pub.pem file not found on the server."
 
-# --- Cryptographic Constants ---
-AES_KEY_SIZE_CRYPTO = 32
-AES_GCM_NONCE_SIZE_CRYPTO = 12
-
-# ==============================================================================
-#  CORE DECRYPTION FUNCTION
-# ==============================================================================
+# --- Core Decryption Function ---
 def decode_data(image_base64_string: str) -> dict:
     """
     This is the core logic function that will be exposed as an API.
@@ -43,15 +37,12 @@ def decode_data(image_base64_string: str) -> dict:
     """
     if not KEYLOCK_PRIV_KEY:
         raise gr.Error("Server Error: The API is not configured with a private key.")
-        
     try:
-        # 1. Decode base64 and load image
         image_buffer = base64.b64decode(image_base64_string)
         private_key = serialization.load_pem_private_key(KEYLOCK_PRIV_KEY.encode(), password=None)
         img = Image.open(io.BytesIO(image_buffer)).convert("RGB")
         pixel_data = np.array(img).ravel()
 
-        # 2. Extract LSB data
         HEADER_BITS = 32
         if pixel_data.size < HEADER_BITS: raise ValueError("Image too small for header.")
         header_binary_string = "".join(str(pixel & 1) for pixel in pixel_data[:HEADER_BITS])
@@ -63,25 +54,22 @@ def decode_data(image_base64_string: str) -> dict:
         data_binary_string = "".join(str(pixel & 1) for pixel in pixel_data[HEADER_BITS:end_offset])
         crypto_payload = int(data_binary_string, 2).to_bytes(data_length, byteorder='big')
         
-        # 3. Decrypt payload
         rsa_key_size_bytes = private_key.key_size // 8
         offset = 4
         encrypted_aes_key_len = struct.unpack('>I', crypto_payload[:offset])[0]
         if encrypted_aes_key_len != rsa_key_size_bytes: raise ValueError("Key pair mismatch.")
         encrypted_aes_key = crypto_payload[offset:offset + encrypted_aes_key_len]
         offset += encrypted_aes_key_len
-        nonce = crypto_payload[offset:offset + AES_GCM_NONCE_SIZE_CRYPTO]
-        offset += AES_GCM_NONCE_SIZE_CRYPTO
+        nonce = crypto_payload[offset:offset + 12] # AES_GCM_NONCE_SIZE_CRYPTO
+        offset += 12
         ciphertext_with_tag = crypto_payload[offset:]
         recovered_aes_key = private_key.decrypt(
             encrypted_aes_key,
             padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
         )
         decrypted_bytes = AESGCM(recovered_aes_key).decrypt(nonce, ciphertext_with_tag, None)
-        
         logger.info("Successfully decoded an image.")
         return json.loads(decrypted_bytes.decode('utf-8'))
-
     except (ValueError, InvalidTag, TypeError) as e:
         logger.error(f"Decryption failed: {e}")
         raise gr.Error(f"Decryption failed. The image may be corrupt, or was not encrypted with the correct public key. Details: {e}")
@@ -94,7 +82,7 @@ def decode_data(image_base64_string: str) -> dict:
 # ==============================================================================
 with gr.Blocks(title="Secure Decoder API") as demo:
     gr.Markdown("# Secure KeyLock Decoder API (Server)")
-    gr.Markdown("This Gradio Space provides the secure, server-side API for decoding KeyLock images. See documentation below.")
+    gr.Markdown("This Gradio Space provides the secure, server-side API for decoding KeyLock images.")
 
     with gr.Tabs():
         with gr.TabItem("API Status & Docs"):
@@ -105,23 +93,24 @@ with gr.Blocks(title="Secure Decoder API") as demo:
             gr.Markdown("---")
             gr.Markdown("## API Documentation")
             gr.Markdown("### Public Key\nUse the following public key to encrypt data for this service.")
-            gr.Code(value=PUBLIC_KEY_PEM_STRING, language="python", label="Service Public Key")
+            gr.Code(value=PUBLIC_KEY_PEM_STRING, language="pem", label="Service Public Key")
             
-            gr.Markdown("### API Endpoint: `/run/mcp_decode` (Base64)")
+            # --- UPDATED DOCUMENTATION TO REFLECT NEW API NAME ---
+            gr.Markdown("### API Endpoint: `/run/keylock-auth-decoder`")
             gr.Markdown("- **Method**: `POST`\n- **Body**: `JSON`")
             gr.Code(language="json", label="JSON Payload Schema", value='{\n  "data": [\n    "iVBORw0KGgoAAAANSUhEUgA..."\n  ]\n}')
-            gr.Code(language="python", label="Example using cURL",
+            gr.Code(language="bash", label="Example using cURL",
                     value="""# Get base64 string: B64_STRING=$(base64 -w 0 /path/to/image.png)
 curl -X POST -H "Content-Type: application/json" \\
 -d '{"data": ["'$B64_STRING'"]}' \\
-https://your-space-name.hf.space/run/mcp_decode""")
+https://your-space-name.hf.space/run/keylock-auth-decoder""")
 
     # --- This is the hidden component that creates the API endpoint ---
-    # We define the function it calls (`decode_data`) and give it an `api_name`.
     with gr.Row(visible=False):
         api_input = gr.Textbox()
         api_output = gr.JSON()
-        gr.Interface(fn=decode_data, inputs=api_input, outputs=api_output, api_name="mcp_decode")
+        # --- RENAMED API ENDPOINT AS REQUESTED ---
+        gr.Interface(fn=decode_data, inputs=api_input, outputs=api_output, api_name="keylock-auth-decoder")
 
 if __name__ == "__main__":
     demo.launch()