Update app.py

app.py

@@ -1,28 +1,106 @@
 """
 Secure KeyLock Decoder API Server
 
-This script deploys a secure Gradio application that acts as a server-side API 
+This script deploys a secure Gradio application that acts as a server-side API
 for decrypting and retrieving JSON data hidden within PNG images.
 
-It implements a hybrid security model:
-1.  **Steganography (LSB):** The encrypted payload is hidden in the least significant 
-    bits (LSB) of the image's pixel data. This makes the data-carrying image 
-    visually indistinguishable from the original.
-2.  **Hybrid Encryption (RSA-KEM + AES-GCM):**
-    - A one-time, symmetric AES-GCM key is generated to encrypt the actual JSON payload.
-    - This AES key is then encrypted using the server's public RSA key.
-    - The final payload embedded in the image consists of the encrypted AES key, 
-      the AES nonce, and the encrypted JSON data (ciphertext).
-
-This server holds the private RSA key required to decrypt the AES key, and subsequently,
-the final payload. The public RSA key must be distributed to clients who wish to
-encrypt data for this service.
-
-Required Environment Variables:
-- `KEYLOCK_PRIV_KEY`: A string containing the server's PEM-encoded private RSA key.
-
-Required Files:
-- `keylock_pub.pem`: A file containing the server's corresponding PEM-encoded public key.
+================================================================================
+▶️ DEPLOYMENT GUIDE
+================================================================================
+
+---
+OPTION 1: DEPLOY TO HUGGING FACE SPACES (RECOMMENDED)
+---
+This is the easiest and most secure way to deploy this application.
+
+1.  **Generate RSA Keys:**
+    First, you need a private/public RSA key pair. Use OpenSSL on your local machine:
+    ```bash
+    # Generate a 4096-bit private key (stronger)
+    openssl genpkey -algorithm RSA -out keylock_priv.pem -pkeyopt rsa_keygen_bits:4096
+
+    # Extract the public key from the private key
+    openssl rsa -pubout -in keylock_priv.pem -out keylock_pub.pem
+    ```
+    This will create two files: `keylock_priv.pem` (keep this secret!) and `keylock_pub.pem` (this is safe to share).
+
+2.  **Create a Hugging Face Space:**
+    - Go to Hugging Face and create a new "Space".
+    - Choose the "Gradio" SDK.
+    - Give it a name (e.g., "my-keylock-decoder").
+
+3.  **Upload Files to the Space Repository:**
+    - Rename this script to `app.py`.
+    - Create a `requirements.txt` file with the following content:
+      ```
+      gradio
+      numpy
+      Pillow
+      cryptography
+      ```
+    - Upload `app.py`, `requirements.txt`, and the public key `keylock_pub.pem` to your Space's repository.
+    - **DO NOT UPLOAD THE PRIVATE KEY (`keylock_priv.pem`)!**
+
+4.  **Set the Private Key as a Secret:**
+    - In your Space, go to the "Settings" tab.
+    - Find the "Repository secrets" section.
+    - Click "New secret".
+    - **Name:** `KEYLOCK_PRIV_KEY` (this name must be exact).
+    - **Value:** Open `keylock_priv.pem` on your local machine, copy its ENTIRE content (including `-----BEGIN PRIVATE KEY-----` and `-----END PRIVATE KEY-----`), and paste it into the value field.
+    - The application will now automatically and securely load this key at runtime.
+
+---
+OPTION 2: RUN LOCALLY FOR DEVELOPMENT
+---
+Use this for testing on your own computer.
+
+1.  **Generate Keys:** Follow Step 1 from the Hugging Face guide.
+
+2.  **Install Dependencies:**
+    ```bash
+    pip install gradio numpy Pillow cryptography
+    ```
+
+3.  **Set Environment Variable:**
+    You must provide the private key as an environment variable.
+    - Open `keylock_priv.pem`, copy its entire content into your clipboard.
+    - In your terminal (Linux/macOS):
+      ```bash
+      export KEYLOCK_PRIV_KEY='PASTE_THE_ENTIRE_KEY_CONTENT_HERE'
+      python app.py
+      ```
+    - In Windows PowerShell:
+      ```powershell
+      $env:KEYLOCK_PRIV_KEY='PASTE_THE_ENTIRE_KEY_CONTENT_HERE'
+      python app.py
+      ```
+
+4.  **Run the Script:** The app will be available at `http://127.0.0.1:7860`.
+
+---
+OPTION 3: DEPLOY TO A SELF-HOSTED SERVER
+---
+For advanced users deploying on their own VPS or server.
+
+1.  **Generate Keys & Install Dependencies:** Follow steps 1 & 2 from the local guide.
+
+2.  **Launch the App:**
+    Modify the `demo.launch()` line at the bottom of this script to bind to all network interfaces:
+    `demo.launch(server_name="0.0.0.0", server_port=7860)`
+
+3.  **Manage Environment Variable:**
+    Set the `KEYLOCK_PRIV_KEY` environment variable using a production-safe method like a `.env` file with `python-dotenv`, systemd service files, or your container orchestration platform (e.g., Docker, Kubernetes).
+
+4.  **Use a Reverse Proxy (CRITICAL):**
+    Do not expose the Gradio port directly to the internet. Place the application behind a reverse proxy like Nginx or Caddy. The proxy will handle SSL/TLS termination (HTTPS), provide better security, and manage traffic.
+
+================================================================================
+
+This application implements a hybrid security model:
+1.  **Steganography (LSB):** The encrypted payload is hidden in the least significant
+    bits (LSB) of the image's pixel data.
+2.  **Hybrid Encryption (RSA-KEM + AES-GCM):** The actual JSON payload is encrypted
+    with a one-time AES key, which itself is encrypted with the server's RSA public key.
 """
 
 import os
@@ -61,7 +139,7 @@ try:
     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."
+    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}"
@@ -90,21 +168,12 @@ def decode_data(image_base64_string: str) -> dict:
     ------------------------------------------------------------------------------------------
     ... | AES Nonce (12 bytes) | AES-GCM Ciphertext + Tag (remaining bytes) |
     ------------------------------------------------------------------------
-
-    The process is as follows:
-    1.  Read the first 32 LSBs to get the total length of the crypto payload.
-    2.  Extract the entire crypto payload from the image's LSBs.
-    3.  Unpack the crypto payload:
-        a. Decrypt the RSA-encrypted AES key using the server's private key.
-        b. Extract the nonce.
-        c. Use the recovered AES key and nonce to decrypt the final ciphertext.
-    4.  Decode the decrypted bytes (UTF-8) and parse them as JSON.
     """
     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)
@@ -114,11 +183,10 @@ def decode_data(image_base64_string: str) -> dict:
         # 2. Extract data length from the LSB header
         if pixel_data.size < HEADER_BITS:
             raise ValueError(f"Image is too small to contain a valid header. 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 length is 0, it's a valid image with no data.
         if data_length == 0:
             logger.info("Image decoded with zero data length. Returning empty dict.")
             return {}
@@ -128,42 +196,36 @@ def decode_data(image_base64_string: str) -> dict:
         end_offset = HEADER_BITS + data_bits_count
         if pixel_data.size < end_offset:
             raise ValueError("Image data corrupt or truncated. The actual image size is smaller than specified in the header.")
-            
+
         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 the crypto payload
         private_key = serialization.load_pem_private_key(KEYLOCK_PRIV_KEY.encode(), password=None)
-        
-        # Unpack Encrypted AES Key
+
         offset = 4 # First 4 bytes define the length of the encrypted AES key
         encrypted_aes_key_len = struct.unpack('>I', crypto_payload[:offset])[0]
-        
-        # Sanity check: ensure the key length matches the server's RSA key size
+
         rsa_key_size_bytes = private_key.key_size // 8
         if encrypted_aes_key_len != rsa_key_size_bytes:
             raise ValueError(f"Key mismatch: The data was encrypted with a {encrypted_aes_key_len*8}-bit key, but the server uses a {private_key.key_size}-bit key.")
 
         encrypted_aes_key = crypto_payload[offset : offset + encrypted_aes_key_len]
         offset += encrypted_aes_key_len
-        
-        # Unpack Nonce
+
         nonce = crypto_payload[offset : offset + AES_GCM_NONCE_SIZE]
         offset += AES_GCM_NONCE_SIZE
-        
-        # The rest is the ciphertext + authentication tag
+
         ciphertext_with_tag = crypto_payload[offset:]
 
         # 5. Decrypt the data
-        # a. Decrypt the AES key with the private RSA key
         recovered_aes_key = private_key.decrypt(
             encrypted_aes_key,
             padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
         )
-        # b. Decrypt the payload with the recovered AES key
         aesgcm = AESGCM(recovered_aes_key)
         decrypted_bytes = aesgcm.decrypt(nonce, ciphertext_with_tag, None)
-        
+
         logger.info("Successfully decoded and decrypted a payload from an image.")
         return json.loads(decrypted_bytes.decode('utf-8'))
 
@@ -179,31 +241,28 @@ def decode_data(image_base64_string: str) -> dict:
 # ==============================================================================
 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.")
+    gr.Markdown("This application provides a secure API endpoint to decrypt and extract JSON data hidden within PNG images. **See the deployment guide at the top of `app.py` for setup instructions.**")
 
     with gr.Tabs():
         with gr.TabItem("🚀 Quick Start & Documentation"):
-            
+
             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", open=False):
+            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. This key will only be used for this one message.
-                    2.  **Client Encrypts Data**: The client takes the JSON payload, converts it to bytes, and encrypts it using the one-time AES key. This produces the `ciphertext` and an authentication `tag`.
-                    3.  **Client Encrypts the One-Time Key**: To send the AES key securely, the client encrypts it using this server's **Public RSA Key** (provided below). Only the server, with its matching private key, can decrypt this.
-                    4.  **Client Creates the Payload**: The client bundles everything into a single binary payload in a specific order:
-                        - The encrypted AES key.
-                        - The `nonce` (a random value used in AES encryption).
-                        - The `ciphertext` and `tag`.
-                    5.  **Client Hides the Payload**: The client uses **Least Significant Bit (LSB) steganography** to hide the binary payload in the pixels of a template PNG image. It also writes a 32-bit header specifying the payload's total length.
-                    6.  **Client Sends Image**: The final image, which looks normal to the human eye, is sent to this API endpoint.
-                    7.  **Server Decodes**: This server reverses the process: it extracts the binary payload from the LSBs, uses its private RSA key to decrypt the one-time AES key, and then uses that key to decrypt the final message.
+                    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.
                     """
                 )
 
@@ -211,70 +270,36 @@ with gr.Blocks(title="Secure Decoder API", theme=gr.themes.Soft()) as demo:
             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="python", label="Server Public Key")
-            
+            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. This code performs the inverse of the server's operation.")
+                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
+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
-from cryptography.hazmat.primitives import serialization
+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):
-    \"\"\"
-    Encrypts a JSON payload and hides it in a template image using hybrid encryption.
-    \"\"\"
-    # 1. Load the server's public key
     public_key = serialization.load_pem_public_key(public_key_pem.encode())
-    
-    # 2. Prepare the data and generate a one-time AES key
     data_to_encrypt = json.dumps(json_data).encode('utf-8')
     aes_key = AESGCM.generate_key(bit_length=256)
-    aesgcm = AESGCM(aes_key)
-    nonce = os.urandom(12) # 96-bit nonce
-    
-    # 3. Encrypt the data with AES-GCM
-    ciphertext_with_tag = aesgcm.encrypt(nonce, data_to_encrypt, None)
-    
-    # 4. Encrypt the one-time AES key with the server's public RSA key
-    encrypted_aes_key = public_key.encrypt(
-        aes_key,
-        padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
-    )
-    
-    # 5. Assemble the final crypto payload
-    # Format: [length of encrypted_aes_key (4 bytes)] [encrypted_aes_key] [nonce (12 bytes)] [ciphertext]
+    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
-    
-    # 6. Hide the payload in the image using LSB steganography
     img = Image.open(template_image_path).convert("RGB")
     pixel_data = np.array(img)
-    
-    payload_binary_string = ''.join(f'{byte:08b}' for byte in payload)
-    header_binary_string = f'{len(payload):032b}' # 32-bit header for payload length
-    full_binary_string = header_binary_string + payload_binary_string
-    
+    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 for the template image. Needs {len(full_binary_string)} pixels, image has {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)):
-        pixel_val = flat_pixels[i]
-        bit_val = int(full_binary_string[i])
-        # Modify the least significant bit
-        flat_pixels[i] = (pixel_val & 0b11111110) | bit_val
-        
-    new_img = Image.fromarray(pixel_data)
-    new_img.save(output_path, "PNG")
+        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 ---
@@ -282,21 +307,9 @@ 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-ghijkl-789012",
-        "permissions": ["read", "write"]
-    }
-
-    # Ensure you have a 'template.png' in the same directory
-    create_keylock_image(
-        public_key_pem=SERVER_PUBLIC_KEY,
-        json_data=my_secret_data,
-        template_image_path="template.png",
-        output_path="encrypted_image.png"
-    )
-                    """
+    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")
@@ -306,7 +319,7 @@ YOUR_SERVER_PUBLIC_KEY_HERE
                 "- **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.Accordion("Python `requests` Example", open=True):
                 gr.Code(
                     language="python",
@@ -315,67 +328,34 @@ YOUR_SERVER_PUBLIC_KEY_HERE
 import requests
 import base64
 
-# Assumes you have created 'encrypted_image.png' using the script from Step 2
 IMAGE_PATH = "encrypted_image.png"
 API_URL = "https://YOUR-SPACE-NAME.hf.space/run/keylock-auth-decoder" # <-- Change this URL
 
-# Read image and encode it to base64
 with open(IMAGE_PATH, "rb") as f:
     b64_string = base64.b64encode(f.read()).decode('utf-8')
 
-# Prepare the JSON payload
-payload = {
-    "data": [b64_string]
-}
-
-# Send the request
-response = requests.post(API_URL, json=payload)
+response = requests.post(API_URL, json={"data": [b64_string]})
 
 if response.status_code == 200:
-    result = response.json()
-    print("Success! Decrypted data:")
-    print(result['data'][0])
+    print("Success! Decrypted data:", response.json()['data'][0])
 else:
-    print(f"Error: {response.status_code}")
-    print(response.text)
+    print(f"Error: {response.status_code}", response.text)
 """
                 )
 
-            with gr.Accordion("cURL Example", open=False):
-                gr.Code(language="python", label="Example using cURL",
-                        value="""
-# 1. Encode your image to base64 and store it in a variable
-# On macOS: B64_STRING=$(base64 /path/to/encrypted_image.png)
-# On Linux: B64_STRING=$(base64 -w 0 /path/to/encrypted_image.png)
-
-# 2. Call the API
-curl -X POST \\
-     -H "Content-Type: application/json" \\
-     -d '{"data": ["'$B64_STRING'"]}' \\
-     https://YOUR-SPACE-NAME.hf.space/run/keylock-auth-decoder
-"""
-                )
-        
         with gr.TabItem("⚙️ Server Status & Error Guide"):
             gr.Markdown("## Server Status")
-            key_status = "✅ Loaded successfully from secrets." if KEYLOCK_PRIV_KEY else "❌ NOT FOUND. The API is non-functional. The administrator must set the `KEYLOCK_PRIV_KEY` secret in the Space settings."
-            gr.Textbox(label="Private Key Status", value=key_status, interactive=False, lines=2)
-            
+            key_status = "✅ Loaded successfully from secrets." if KEYLOCK_PRIV_KEY else "❌ NOT FOUND. The API is non-functional. The administrator must set the `KEYLOCK_PRIV_KEY` secret in the Space settings. See the deployment guide in `app.py`."
+            gr.Textbox(label="Private Key Status", value=key_status, interactive=False, lines=3)
+
             gr.Markdown("---")
             gr.Markdown("## Error Handling Guide")
             gr.Markdown(
                 """
-                If your API call fails, check the error message for clues:
-
-                - **`"Decryption failed. ... InvalidTag"`**: This is the most common error. It almost always means one of two things:
-                    1. The image was encrypted with a **different public key** than the one this server uses.
-                    2. The image file was corrupted during transfer or storage, altering the hidden data.
-
-                - **`"Decryption failed. ... Key mismatch"`**: The RSA key size used to encrypt the data (e.g., 2048-bit) does not match the server's key size (e.g., 4096-bit). Ensure you are using the exact public key provided by this service.
-                
-                - **`"Image data corrupt or truncated..."`**: The header inside the image indicates a certain data length, but the image is not large enough to contain that much data. The file is likely incomplete.
-
-                - **`"Server Error: ... not configured"`**: This is a server-side problem. The administrator has not correctly configured the private key secret for this Gradio Space.
+                - **`"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.
                 """
             )
 
@@ -384,13 +364,13 @@ curl -X POST \\
         api_input = gr.Textbox()
         api_output = gr.JSON()
         gr.Interface(
-            fn=decode_data, 
-            inputs=api_input, 
-            outputs=api_output, 
+            fn=decode_data,
+            inputs=api_input,
+            outputs=api_output,
             api_name="keylock-auth-decoder",
-            # No need for title/description here as it's a hidden API generator
         )
 
 if __name__ == "__main__":
-    # Launch the Gradio app
+    # To run on a server, you might use: demo.launch(server_name="0.0.0.0")
+    # See the deployment guide at the top of this file for more details.
     demo.launch()