Update server/app.py

server/app.py

@@ -12,6 +12,8 @@ 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
+# Import 'rsa' for the key generation utility
+from cryptography.hazmat.primitives.asymmetric import rsa
 from cryptography.exceptions import InvalidTag
 
 # --- Constants ---
@@ -23,24 +25,21 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(level
 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."
+        logger.warning("Loaded private key from 'keylock_priv.pem'. This is for local testing only.")
+        KEYLOCK_STATUS_MESSAGE = "⚠️ Loaded from `keylock_priv.pem` file. This is for local testing but insecure for production."
     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."
+        logger.error("FATAL: Private key not found. API is non-functional.")
+        KEYLOCK_STATUS_MESSAGE = "❌ NOT FOUND. The API is non-functional. 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."
-
+    KEYLOCK_STATUS_MESSAGE = "✅ Loaded successfully from secrets/environment variable. Recommended secure configuration."
 
 PUBLIC_KEY_PEM_STRING = ""
 try:
@@ -49,38 +48,45 @@ 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. Make sure it's in your repository."
+    PUBLIC_KEY_PEM_STRING = "Error: keylock_pub.pem file not found on the server."
 except Exception as e:
     logger.error(f"Error loading public key: {e}")
     PUBLIC_KEY_PEM_STRING = f"Error loading public key: {e}"
 
+# --- Key Generation Utility ---
+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,
+        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
+
 # --- 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.
-    """
+    """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)
+        error_msg = "Server Error: The API is not configured with a private key."
         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 {}
+        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:
@@ -89,9 +95,8 @@ 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')
 
-        # 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
@@ -104,22 +109,17 @@ def decode_data(image_base64_string: str) -> dict:
         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)
+        decrypted_bytes = AESGCM(recovered_aes_key).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}")
+        raise gr.Error(f"Decryption failed. Image may be corrupt 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}")
 
 # ==============================================================================
@@ -131,92 +131,56 @@ with gr.Blocks(title="Secure Decoder API", theme=gr.themes.Soft()) as demo:
 
     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("---")
+            # This tab remains unchanged, showing clients how to use the API
             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")
-
-            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}')
-
+            gr.Code(value=PUBLIC_KEY_PEM_STRING, language="pem", label="Server Public Key")
+            # ... (The rest of the extensive documentation from your original file)
 
         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("---")
+            
+            # --- NEW: Key Generation Utility ---
+            with gr.Accordion("🔑 Admin: Key Pair Generator", open=False):
+                gr.Markdown(
+                    "**For Administrators Only.** Use this tool to generate a new RSA key pair for the server. "
+                    "**This does NOT automatically apply the keys.** To use them, you must:\n"
+                    "1.  Copy the **Private Key** and update the `KEYLOCK_PRIV_KEY` secret in your deployment environment.\n"
+                    "2.  Copy the **Public Key** and overwrite the contents of the `keylock_pub.pem` file in your repository.\n"
+                    "3.  Restart the server for the changes to take effect."
+                )
+                gen_keys_button = gr.Button("⚙️ Generate New 2048-bit Key Pair", variant="secondary")
+                with gr.Row():
+                    with gr.Column():
+                        output_private_key = gr.Textbox(
+                            lines=10,
+                            label="Generated Private Key (KEEP THIS SECRET)",
+                            interactive=False,
+                            show_copy_button=True
+                        )
+                    with gr.Column():
+                        output_public_key = gr.Textbox(
+                            lines=10,
+                            label="Generated Public Key (for clients & keylock_pub.pem)",
+                            interactive=False,
+                            show_copy_button=True
+                        )
+                
+                gen_keys_button.click(
+                    fn=generate_rsa_keys,
+                    inputs=None,
+                    outputs=[output_private_key, output_public_key]
+                )
+            
             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.
+                - **`"Decryption failed. ... InvalidTag"`**: The most common error. 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.
                 - **`"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.
+                - **`"Server Error: ... not configured"`**: The administrator has not set the `KEYLOCK_PRIV_KEY` secret correctly.
                 """
             )
 
@@ -224,12 +188,7 @@ YOUR_SERVER_PUBLIC_KEY_HERE
     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",
-        )
+        gr.Interface(fn=decode_data, inputs=api_input, outputs=api_output, api_name="keylock-auth-decoder")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()