Update server/app.py

server/app.py

@@ -11,60 +11,80 @@ 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.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.primitives.asymmetric import padding, rsa
 from cryptography.exceptions import InvalidTag
 
+# --- Constants ---
 HEADER_BITS = 32
 AES_GCM_NONCE_SIZE = 12
 
+# --- Setup Logging ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-KEYLOCK_PRIV_KEY = os.environ.get('KEYLOCK_PRIV_KEY','./DEMO_ONLY_THIS_IS_SECRET_keylock_priv_key.pem')
+# --- Key Loading and Management ---
+KEYLOCK_PRIV_KEY_PEM = os.environ.get('KEYLOCK_PRIV_KEY')
+PRIVATE_KEY_OBJECT = None
+PUBLIC_KEY_PEM_STRING = ""
 KEYLOCK_STATUS_MESSAGE = ""
-KEY_DIR = "./keys/"
-if not KEYLOCK_PRIV_KEY:
+
+# Prioritize environment variable, then fallback to local file for development
+if not KEYLOCK_PRIV_KEY_PEM:
+    # This path is relative to server/app.py, assuming the Space structure is `server/keys/`
+    dev_key_path = 'keys/DEMO_ONLY_THIS _IS_SECRET_keylock_priv_key.pem'
     try:
-        with open(f"{KEY_DIR}keylock_priv.pem", "r") as f:
-            KEYLOCK_PRIV_KEY = f.read()
-        logger.warning(f"Loaded private key from '{KEY_DIR}keylock_priv.pem'. This is for local testing only.")
-        KEYLOCK_STATUS_MESSAGE = f"⚠️ Loaded from `{KEY_DIR}keylock_priv.pem` file. This is for local testing but insecure for production."
+        with open(dev_key_path, "r") as f:
+            KEYLOCK_PRIV_KEY_PEM = f.read()
+        logger.warning(f"Loaded private key from '{dev_key_path}'. This is for local testing only.")
+        KEYLOCK_STATUS_MESSAGE = f"⚠️ Loaded from `{dev_key_path}`. This is for local testing but insecure for production."
     except FileNotFoundError:
-        logger.error("FATAL: Private key not found. API is non-functional.")
-        KEYLOCK_STATUS_MESSAGE = f"❌ NOT FOUND. The API is non-functional. Set the `KEYLOCK_PRIV_KEY` secret or provide a `{KEY_DIR}keylock_priv.pem` file."
+        logger.error(f"FATAL: Private key not found at '{dev_key_path}' and 'KEYLOCK_PRIV_KEY' secret is not set.")
+        KEYLOCK_STATUS_MESSAGE = "❌ NOT FOUND. The API is non-functional. Set the `KEYLOCK_PRIV_KEY` secret or provide the demo key file."
 else:
     logger.info("Successfully loaded private key from environment variable 'KEYLOCK_PRIV_KEY'.")
     KEYLOCK_STATUS_MESSAGE = "✅ Loaded successfully from secrets/environment variable. 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."
-except Exception as e:
-    logger.error(f"Error loading public key: {e}")
-    PUBLIC_KEY_PEM_STRING = f"Error loading public key: {e}"
-
-def generate_rsa_keys():
-    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
+# Derive public key from private key if available
+if KEYLOCK_PRIV_KEY_PEM:
+    try:
+        PRIVATE_KEY_OBJECT = serialization.load_pem_private_key(KEYLOCK_PRIV_KEY_PEM.encode(), password=None)
+        public_key = PRIVATE_KEY_OBJECT.public_key()
+        PUBLIC_KEY_PEM_STRING = public_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo
+        ).decode('utf-8')
+        KEYLOCK_STATUS_MESSAGE += "\n✅ Public key derived successfully."
+    except Exception as e:
+        logger.error(f"Failed to load private key or derive public key: {e}", exc_info=True)
+        PRIVATE_KEY_OBJECT = None
+        PUBLIC_KEY_PEM_STRING = f"Error: Could not process the configured private key. Details: {e}"
+        KEYLOCK_STATUS_MESSAGE += f"\n❌ Failed to parse key: {e}"
+
+# --- API Functions ---
+
+def get_public_key():
+    """API endpoint to return the server's public key."""
+    if not PUBLIC_KEY_PEM_STRING or "Error" in PUBLIC_KEY_PEM_STRING:
+        raise gr.Error("Server key is not configured correctly.")
+    return PUBLIC_KEY_PEM_STRING
+
+def get_server_info():
+    """API endpoint to return server documentation and metadata."""
+    return {
+        "name": "KeyLock Auth Server",
+        "version": "1.2",
+        "documentation": "This server decrypts data hidden in KeyLock images. Use the /keylock-pub endpoint to get the public key for encryption, and POST to /keylock-server with a base64 encoded image string to decrypt.",
+        "endpoints": {
+            "/keylock-pub": "GET - Returns the server's public key.",
+            "/keylock-info": "GET - Returns this information object.",
+            "/keylock-server": "POST - Decrypts a KeyLock image."
+        }
+    }
 
 def decode_data(image_base64_string: str) -> dict:
-    if not KEYLOCK_PRIV_KEY:
+    if not PRIVATE_KEY_OBJECT:
         error_msg = "Server Error: The API is not configured with a private key."
+        logger.error(error_msg)
         raise gr.Error(error_msg)
     try:
         image_buffer = base64.b64decode(image_base64_string)
@@ -87,13 +107,11 @@ 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')
 
-        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
+        rsa_key_size_bytes = PRIVATE_KEY_OBJECT.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.")
+            raise ValueError(f"Key mismatch: Encrypted with a {encrypted_aes_key_len*8}-bit key, server uses {PRIVATE_KEY_OBJECT.key_size}-bit key.")
 
         encrypted_aes_key = crypto_payload[offset : offset + encrypted_aes_key_len]
         offset += encrypted_aes_key_len
@@ -101,7 +119,7 @@ def decode_data(image_base64_string: str) -> dict:
         offset += AES_GCM_NONCE_SIZE
         ciphertext_with_tag = crypto_payload[offset:]
 
-        recovered_aes_key = private_key.decrypt(
+        recovered_aes_key = PRIVATE_KEY_OBJECT.decrypt(
             encrypted_aes_key,
             padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
         )
@@ -110,100 +128,79 @@ def decode_data(image_base64_string: str) -> dict:
         return json.loads(decrypted_bytes.decode('utf-8'))
 
     except (ValueError, InvalidTag, TypeError, struct.error) as e:
+        logger.warning(f"Decryption failed: {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 server error occurred during decryption: {e}", exc_info=True)
         raise gr.Error(f"An unexpected server error occurred. Details: {e}")
 
+def generate_rsa_keys():
+    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
+
+# --- Gradio UI ---
 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 secure API endpoints to decrypt and extract JSON data hidden within PNG images.")
 
     with gr.Tabs():
-        with gr.TabItem("🚀 Quick Start & Documentation"):
+        with gr.TabItem("🚀 API Documentation"):
             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")
-            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.Markdown(
+                "This server exposes three main API endpoints for programmatic use:\n"
+                "1.  **`/keylock-info`**: A `GET` request returns a JSON object with server metadata.\n"
+                "2.  **`/keylock-pub`**: A `GET` request returns the server's public RSA key required for encryption.\n"
+                "3.  **`/keylock-server`**: A `POST` request with a base64-encoded image string decrypts the data."
+            )
+            gr.Markdown("### Server's Public Key")
+            gr.Code(value=PUBLIC_KEY_PEM_STRING, language="pem", label="Server Public Key (from /keylock-pub)")
+            with gr.Accordion("Client-Side Python Example for Image Creation", open=False):
                 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")
+# See the KeyLock-Auth-Client repo for a full implementation.
+# The core logic involves using the server's public key to encrypt
+# a one-time AES key, which in turn encrypts your JSON payload.
+# This encrypted bundle is then embedded into an image's pixel data.
 """
                 )
 
-        with gr.TabItem("⚙️ Server Status & Error Guide"):
+        with gr.TabItem("⚙️ Server Status & Admin"):
             gr.Markdown("## Server Status")
             gr.Textbox(label="Private Key Status", value=KEYLOCK_STATUS_MESSAGE, interactive=False, lines=3)
             gr.Markdown("---")
-            
+
             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."
+                    "2.  Restart the server for the changes to take effect. The public key will be derived automatically."
                 )
                 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. 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.
-                """
-            )
+                    output_private_key = gr.Textbox(lines=10, label="Generated Private Key (KEEP THIS SECRET)", interactive=False, show_copy_button=True)
+                    output_public_key = gr.Textbox(lines=10, label="Generated Public Key (will be auto-derived)", interactive=False, show_copy_button=True)
+                gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
 
+    # --- API Endpoints (Hidden from UI) ---
     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")
+        # API Endpoint 1: Public Key
+        gr.Interface(fn=get_public_key, inputs=None, outputs=gr.Textbox(), api_name="keylock-pub")
+        # API Endpoint 2: Info
+        gr.Interface(fn=get_server_info, inputs=None, outputs=gr.JSON(), api_name="keylock-info")
+        # API Endpoint 3: Decrypt (renamed)
+        gr.Interface(fn=decode_data, inputs=gr.Textbox(), outputs=gr.JSON(), api_name="keylock-server")
 
 if __name__ == "__main__":
     demo.launch()