Create server/app.py

server/app.py

@@ -0,0 +1,235 @@
+import os
+import io
+import json
+import base64
+import struct
+import logging
+
+import gradio as gr
+from PIL import Image
+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
+
+# --- Constants ---
+HEADER_BITS = 32
+AES_GCM_NONCE_SIZE = 12
+
+# --- Configure Logging ---
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# --- Load Keys and Config ---
+# For production (e.g., Hugging Face Spaces), the private key should be an environment variable.
+KEYLOCK_PRIV_KEY = os.environ.get('KEYLOCK_PRIV_KEY')
+KEYLOCK_STATUS_MESSAGE = ""
+
+# For local testing, fall back to reading from a file if the environment variable is not set.
+if not KEYLOCK_PRIV_KEY:
+    try:
+        with open("keylock_priv.pem", "r") as f:
+            KEYLOCK_PRIV_KEY = f.read()
+        logger.warning("Loaded private key from 'keylock_priv.pem'. This is for local testing only. Use environment variables in production.")
+        KEYLOCK_STATUS_MESSAGE = "⚠️ Loaded from `keylock_priv.pem` file. This is acceptable for local testing but insecure for production. Use secrets/environment variables for deployment."
+    except FileNotFoundError:
+        logger.error("FATAL: Private key not found as env var 'KEYLOCK_PRIV_KEY' or in 'keylock_priv.pem'. API is non-functional.")
+        KEYLOCK_STATUS_MESSAGE = "❌ NOT FOUND. The API is non-functional. The administrator must set the `KEYLOCK_PRIV_KEY` secret or provide a `keylock_priv.pem` file."
+else:
+    logger.info("Successfully loaded private key from environment variable 'KEYLOCK_PRIV_KEY'.")
+    KEYLOCK_STATUS_MESSAGE = "✅ Loaded successfully from secrets/environment variable. This is the recommended secure configuration."
+
+
+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 for display.")
+except FileNotFoundError:
+    logger.error("FATAL: keylock_pub.pem not found. The UI will show an error.")
+    PUBLIC_KEY_PEM_STRING = "Error: keylock_pub.pem file not found on the server. Make sure it's in your repository."
+except Exception as e:
+    logger.error(f"Error loading public key: {e}")
+    PUBLIC_KEY_PEM_STRING = f"Error loading public key: {e}"
+
+# --- Core Decryption Function ---
+def decode_data(image_base64_string: str) -> dict:
+    """
+    Decrypts a payload hidden in a PNG image via LSB steganography and hybrid encryption.
+    """
+    if not KEYLOCK_PRIV_KEY:
+        error_msg = "Server Error: The API is not configured with a private key. Please contact the administrator."
+        logger.error(error_msg)
+        raise gr.Error(error_msg)
+
+    try:
+        # 1. Decode Base64 and load image data
+        image_buffer = base64.b64decode(image_base64_string)
+        img = Image.open(io.BytesIO(image_buffer)).convert("RGB")
+        pixel_data = np.array(img).ravel()
+
+        # 2. Extract data length from the LSB header
+        if pixel_data.size < HEADER_BITS:
+            raise ValueError(f"Image is too small. Minimum pixel count: {HEADER_BITS}")
+
+        header_binary_string = "".join(str(pixel & 1) for pixel in pixel_data[:HEADER_BITS])
+        data_length = int(header_binary_string, 2)
+
+        if data_length == 0:
+            return {}
+
+        # 3. Extract the full crypto payload
+        data_bits_count = data_length * 8
+        end_offset = HEADER_BITS + data_bits_count
+        if pixel_data.size < end_offset:
+            raise ValueError("Image data corrupt or truncated.")
+
+        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')
+
+        # 4. Load private key and unpack payload
+        private_key = serialization.load_pem_private_key(KEYLOCK_PRIV_KEY.encode(), password=None)
+
+        offset = 4
+        encrypted_aes_key_len = struct.unpack('>I', crypto_payload[:offset])[0]
+        rsa_key_size_bytes = private_key.key_size // 8
+        if encrypted_aes_key_len != rsa_key_size_bytes:
+            raise ValueError(f"Key mismatch: Encrypted with a {encrypted_aes_key_len*8}-bit key, server uses {private_key.key_size}-bit key.")
+
+        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]
+        offset += AES_GCM_NONCE_SIZE
+        ciphertext_with_tag = crypto_payload[offset:]
+
+        # 5. Decrypt
+        recovered_aes_key = private_key.decrypt(
+            encrypted_aes_key,
+            padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
+        )
+        aesgcm = AESGCM(recovered_aes_key)
+        decrypted_bytes = aesgcm.decrypt(nonce, ciphertext_with_tag, None)
+
+        logger.info("Successfully decrypted payload from image.")
+        return json.loads(decrypted_bytes.decode('utf-8'))
+
+    except (ValueError, InvalidTag, TypeError, struct.error) as e:
+        logger.error(f"Decryption failed: {e}")
+        raise gr.Error(f"Decryption failed. The image may be corrupt, not from a compatible client, or used the wrong public key. Details: {e}")
+    except Exception as e:
+        logger.error(f"An unexpected error occurred: {e}", exc_info=True)
+        raise gr.Error(f"An unexpected server error occurred. Details: {e}")
+
+# ==============================================================================
+#  GRADIO INTERFACE
+# ==============================================================================
+with gr.Blocks(title="Secure Decoder API", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🔐 Secure KeyLock Decoder API")
+    gr.Markdown("This application provides a secure API endpoint to decrypt and extract JSON data hidden within PNG images.")
+
+    with gr.Tabs():
+        with gr.TabItem("🚀 Quick Start & Documentation"):
+            # ... (The rest of your extensive Gradio UI markdown goes here, it's perfect as is)
+            gr.Markdown("## What Is This?")
+            gr.Markdown(
+                "This service decodes messages that have been securely hidden inside images. It's designed for use cases where you need to transmit a small, sensitive JSON payload (like authentication tokens, user data, or configuration) "
+                "without it being easily detectable or readable. The process uses a combination of **steganography** (hiding data in the image's pixels) and **hybrid encryption** (the security of RSA combined with the speed of AES)."
+            )
+
+            with gr.Accordion("How It Works (The Gory Details)", open=False):
+                gr.Markdown(
+                    """
+                    The security relies on a well-established cryptographic pattern called **Hybrid Encryption**. Here’s the step-by-step process a client must follow to create a compatible image:
+                    1.  **Client Generates a One-Time Key**: The client creates a random, single-use 32-byte symmetric key for AES-GCM encryption.
+                    2.  **Client Encrypts Data**: The client takes the JSON payload and encrypts it using the one-time AES key.
+                    3.  **Client Encrypts the One-Time Key**: The client encrypts the AES key using this server's **Public RSA Key** (provided below).
+                    4.  **Client Creates the Payload**: The client bundles the encrypted AES key, nonce, and ciphertext into a single binary payload.
+                    5.  **Client Hides the Payload**: The client uses **Least Significant Bit (LSB) steganography** to hide the payload in a template image.
+                    6.  **Client Sends Image**: The final image is sent to this API endpoint.
+                    7.  **Server Decodes**: This server reverses the process: it extracts the payload, uses its private RSA key to decrypt the AES key, and then decrypts the final message.
+                    """
+                )
+
+            gr.Markdown("---")
+            gr.Markdown("## How to Use This API")
+            gr.Markdown("### Step 1: Get the Server's Public Key")
+            gr.Markdown("You must use this public key to encrypt data intended for this server. Save it to a file (e.g., `server_pub.pem`).")
+            gr.Code(value=PUBLIC_KEY_PEM_STRING, language="pem", label="Server Public Key")
+
+            with gr.Accordion("Step 2: Create the Encrypted Image (Client-Side Python Example)", open=True):
+                gr.Markdown("Use the following Python code in your client application to create a KeyLock image.")
+                gr.Code(
+                    language="python",
+                    label="Client-Side Image Creation Script (create_keylock_image.py)",
+                    value="""
+import json; import os; import struct; from PIL import Image; import numpy as np
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import padding
+def create_keylock_image(public_key_pem: str, json_data: dict, template_image_path: str, output_path: str):
+    public_key = serialization.load_pem_public_key(public_key_pem.encode())
+    data_to_encrypt = json.dumps(json_data).encode('utf-8')
+    aes_key = AESGCM.generate_key(bit_length=256)
+    nonce = os.urandom(12)
+    ciphertext_with_tag = AESGCM(aes_key).encrypt(nonce, data_to_encrypt, None)
+    encrypted_aes_key = public_key.encrypt(aes_key, padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None))
+    payload = struct.pack('>I', len(encrypted_aes_key)) + encrypted_aes_key + nonce + ciphertext_with_tag
+    img = Image.open(template_image_path).convert("RGB")
+    pixel_data = np.array(img)
+    full_binary_string = f'{len(payload):032b}' + ''.join(f'{byte:08b}' for byte in payload)
+    if len(full_binary_string) > pixel_data.size:
+        raise ValueError(f"Data is too large. Needs {len(full_binary_string)} pixels, image has {pixel_data.size}.")
+    flat_pixels = pixel_data.ravel()
+    for i in range(len(full_binary_string)):
+        flat_pixels[i] = (flat_pixels[i] & 0b11111110) | int(full_binary_string[i])
+    Image.fromarray(pixel_data).save(output_path, "PNG")
+    print(f"Successfully created encrypted image at {output_path}")
+# --- Example Usage ---
+if __name__ == '__main__':
+    SERVER_PUBLIC_KEY = \"\"\"-----BEGIN PUBLIC KEY-----
+YOUR_SERVER_PUBLIC_KEY_HERE
+-----END PUBLIC KEY-----\"\"\"
+    my_secret_data = {"user_id": "user-12345", "token": "abcdef-123456"}
+    create_keylock_image(SERVER_PUBLIC_KEY, my_secret_data, "template.png", "encrypted_image.png")
+"""
+                )
+
+            gr.Markdown("### Step 3: Call the API Endpoint")
+            gr.Markdown(
+                "- **Method**: `POST`\n"
+                "- **Endpoint**: `/run/keylock-auth-decoder`\n"
+                "- **Body**: A JSON object containing a base64-encoded string of the image."
+            )
+            gr.Code(language="json", label="JSON Payload Schema", value='{\n  "data": [\n    "<base64_encoded_image_string>"\n  ]\n}')
+
+
+        with gr.TabItem("⚙️ Server Status & Error Guide"):
+            gr.Markdown("## Server Status")
+            gr.Textbox(label="Private Key Status", value=KEYLOCK_STATUS_MESSAGE, interactive=False, lines=3)
+
+            gr.Markdown("---")
+            gr.Markdown("## Error Handling Guide")
+            gr.Markdown(
+                """
+                - **`"Decryption failed. ... InvalidTag"`**: The most common error. It means the image was encrypted with a **different public key** or the image file was corrupted.
+                - **`"Decryption failed. ... Key mismatch"`**: The RSA key size used to encrypt the data does not match the server's key size. Use the public key provided by this service.
+                - **`"Image data corrupt or truncated..."`**: The image file is incomplete or damaged.
+                - **`"Server Error: ... not configured"`**: A server-side problem. The administrator has not set the `KEYLOCK_PRIV_KEY` secret correctly.
+                """
+            )
+
+    # --- Hidden component for the API endpoint ---
+    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="keylock-auth-decoder",
+        )
+
+if __name__ == "__main__":
+    demo.launch()