Update server/app.py

server/app.py

@@ -12,19 +12,15 @@ 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 ---
 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 ---
 KEYLOCK_PRIV_KEY = os.environ.get('KEYLOCK_PRIV_KEY')
 KEYLOCK_STATUS_MESSAGE = ""
 
@@ -53,9 +49,7 @@ 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,
@@ -68,9 +62,7 @@ def generate_rsa_keys():
     ).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."""
     if not KEYLOCK_PRIV_KEY:
         error_msg = "Server Error: The API is not configured with a private key."
         raise gr.Error(error_msg)
@@ -122,27 +114,51 @@ def decode_data(image_base64_string: str) -> dict:
     except Exception as e:
         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"):
-            # 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.Code(value=PUBLIC_KEY_PEM_STRING, language="python", label="Server Public Key")
-            # ... (The rest of the extensive documentation from your original file)
+            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=False):
+                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",
+                    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))
+    crypto_payload = struct.pack('>I', len(encrypted_aes_key)) + encrypted_aes_key + nonce + ciphertext_with_tag
+    header = struct.pack('>I', len(crypto_payload))
+    full_binary_string = ''.join(f'{byte:08b}' for byte in header + crypto_payload)
+    img = Image.open(template_image_path).convert("RGB")
+    pixel_data = np.array(img)
+    if len(full_binary_string) > pixel_data.size:
+        raise ValueError("Data is too large for the image.")
+    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")
+"""
+                )
 
         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. "
@@ -177,14 +193,13 @@ with gr.Blocks(title="Secure Decoder API", theme=gr.themes.Soft()) as demo:
             gr.Markdown("## Error Handling Guide")
             gr.Markdown(
                 """
-                - **`"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"`**: The administrator has not set the `KEYLOCK_PRIV_KEY` secret correctly.
+- **`"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"`**: 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()