Rename server/app.py to server/server.py

server/{app.py → server.py}

@@ -14,29 +14,24 @@ from cryptography.hazmat.primitives import serialization
 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__)
 
-# --- 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 = ""
 
-# 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'
+    dev_key_path = os.path.join(os.path.dirname(__file__), '..', 'keys', 'DEMO_ONLY_THIS _IS_SECRET_keylock_priv_key.pem')
     try:
         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."
+        logger.warning(f"Loaded private key from dev path. This is for local testing only.")
+        KEYLOCK_STATUS_MESSAGE = f"⚠️ Loaded from development key file. This is for local testing but insecure for production."
     except FileNotFoundError:
         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."
@@ -44,7 +39,6 @@ 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."
 
-# 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)
@@ -60,24 +54,36 @@ if KEYLOCK_PRIV_KEY_PEM:
         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 ---
+MOCK_USER_DATABASE = {
+    "sk-12345-abcde": {"user": "demo-user", "permissions": "read"},
+    "sk-67890-fghij": {"user": "admin-user", "permissions": "read,write,delete"}
+}
+
+def example_authenticate(api_key: str, user_id: str) -> bool:
+    if not api_key or not user_id:
+        return False
+    db_entry = MOCK_USER_DATABASE.get(api_key)
+    if db_entry and db_entry.get("user") == user_id:
+        logger.info(f"Authentication successful for user '{user_id}'.")
+        return True
+    else:
+        logger.warning(f"Authentication failed for user '{user_id}' with key '{api_key[:8]}...'.")
+        return False
 
 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.",
+        "version": "1.3",
+        "documentation": "This server decrypts data hidden in KeyLock images and performs a mock authentication. Use /keylock-pub to get the public key, and POST to /keylock-server with a base64 image string to decrypt and authenticate.",
         "endpoints": {
             "/keylock-pub": "GET - Returns the server's public key.",
             "/keylock-info": "GET - Returns this information object.",
-            "/keylock-server": "POST - Decrypts a KeyLock image."
+            "/keylock-server": "POST - Decrypts a KeyLock image and attempts authentication."
         }
     }
 
@@ -90,43 +96,41 @@ def decode_data(image_base64_string: str) -> dict:
         image_buffer = base64.b64decode(image_base64_string)
         img = Image.open(io.BytesIO(image_buffer)).convert("RGB")
         pixel_data = np.array(img).ravel()
-
-        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: raise ValueError("No data found in image.")
         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.")
-
+        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')
-
         offset = 4
         encrypted_aes_key_len = struct.unpack('>I', crypto_payload[:offset])[0]
-        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_OBJECT.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:]
-
-        recovered_aes_key = PRIVATE_KEY_OBJECT.decrypt(
-            encrypted_aes_key,
-            padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
-        )
+        recovered_aes_key = PRIVATE_KEY_OBJECT.decrypt(encrypted_aes_key, padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None))
         decrypted_bytes = AESGCM(recovered_aes_key).decrypt(nonce, ciphertext_with_tag, None)
-
-        return json.loads(decrypted_bytes.decode('utf-8'))
-
+        decrypted_payload = json.loads(decrypted_bytes.decode('utf-8'))
+        logger.info(f"Successfully decoded payload: {decrypted_payload}")
+        api_key = decrypted_payload.get('API_KEY')
+        user_id = decrypted_payload.get('USER')
+        is_authenticated = example_authenticate(api_key=api_key, user_id=user_id)
+        if is_authenticated:
+            return {
+                "authentication_status": "Success",
+                "message": f"User '{user_id}' successfully authenticated.",
+                "granted_permissions": MOCK_USER_DATABASE[api_key]['permissions'],
+                "decoded_payload": decrypted_payload
+            }
+        else:
+            return {
+                "authentication_status": "Failed",
+                "message": "Authentication failed. Invalid credentials provided in the image.",
+                "decoded_payload": decrypted_payload
+            }
     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}")
@@ -136,71 +140,6 @@ def decode_data(image_base64_string: str) -> dict:
 
 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 secure API endpoints to decrypt and extract JSON data hidden within PNG images.")
-
-    with gr.Tabs():
-        with gr.TabItem("🚀 API Documentation"):
-            gr.Markdown("## How to Use This API")
-            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="python", 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="""
-# 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 & 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.  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():
-                    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 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()
+    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