Update app.py

app.py

@@ -13,10 +13,12 @@ from cryptography.hazmat.primitives.asymmetric import rsa, padding
 from cryptography.hazmat.primitives.ciphers.aead import AESGCM
 from cryptography.hazmat.primitives import hashes
 
+# --- Basic Configuration ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-CREATOR_ENDPOINTS_JSON_URL = "https://huggingface.co/spaces/broadfield-dev/KeyLock-Auth-Creator/raw/main/endpoints.json"
+# --- Constants ---
+CREATOR_ENDPOINTS_JSON_URL = "./endpoints.json"
 BASE_HF_URL = "https://huggingface.co/spaces/"
 CREATOR_SPACE_ID = "broadfield-dev/KeyLock-Auth-Creator"
 SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server"
@@ -27,102 +29,180 @@ SERVER_APP_PY_URL = f"{SERVER_URL}/raw/main/app.py"
 
 
 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')
+    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
 
 def create_encrypted_image(secret_data_str: str, public_key_pem: str) -> Image.Image:
-    if not secret_data_str.strip(): raise ValueError("Secret data cannot be empty.")
-    if not public_key_pem.strip(): raise ValueError("Public Key cannot be empty.")
+    """Encrypts data and embeds it into a new image using steganography."""
+    if not secret_data_str.strip():
+        raise ValueError("Secret data cannot be empty.")
+    if not public_key_pem.strip():
+        raise ValueError("Public Key cannot be empty.")
+        
     data_dict = {}
     for line in secret_data_str.splitlines():
         line = line.strip()
-        if not line or line.startswith('#'): continue
+        if not line or line.startswith('#'):
+            continue
         parts = line.split(':', 1) if ':' in line else line.split('=', 1)
-        if len(parts) == 2: data_dict[parts[0].strip()] = parts[1].strip().strip("'\"")
-    if not data_dict: raise ValueError("No valid key-value pairs found.")
+        if len(parts) == 2:
+            data_dict[parts[0].strip()] = parts[1].strip().strip("'\"")
+            
+    if not data_dict:
+        raise ValueError("No valid key-value pairs found in secret data.")
+        
     json_bytes = json.dumps(data_dict).encode('utf-8')
     public_key = serialization.load_pem_public_key(public_key_pem.encode('utf-8'))
+    
+    # AES-GCM encryption
     aes_key, nonce = os.urandom(32), os.urandom(12)
     ciphertext = AESGCM(aes_key).encrypt(nonce, json_bytes, None)
-    rsa_encrypted_key = public_key.encrypt(aes_key, padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None))
+    
+    # RSA-OAEP encryption for the AES key
+    rsa_encrypted_key = public_key.encrypt(
+        aes_key,
+        padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
+    )
+    
+    # Pack the payload: [len(rsa_key)][rsa_key][nonce][ciphertext]
     encrypted_payload = struct.pack('>I', len(rsa_encrypted_key)) + rsa_encrypted_key + nonce + ciphertext
+    
+    # Create a base image
     img = Image.new('RGB', (800, 600), color=(45, 52, 54))
     draw = ImageDraw.Draw(img)
-    try: font = ImageFont.truetype("DejaVuSans.ttf", 40)
-    except IOError: font = ImageFont.load_default(size=30)
+    try:
+        font = ImageFont.truetype("DejaVuSans.ttf", 40)
+    except IOError:
+        font = ImageFont.load_default(size=30)
     draw.text((400, 300), "KeyLock Secure Data", fill=(223, 230, 233), font=font, anchor="ms")
+    
+    # LSB Steganography
     pixel_data = np.array(img.convert("RGB")).ravel()
     binary_payload = ''.join(format(b, '08b') for b in struct.pack('>I', len(encrypted_payload)) + encrypted_payload)
-    if len(binary_payload) > pixel_data.size: raise ValueError("Data too large for image.")
+    
+    if len(binary_payload) > pixel_data.size:
+        raise ValueError(f"Data is too large for the image. Max size: {pixel_data.size // 8} bytes.")
+        
     pixel_data[:len(binary_payload)] = (pixel_data[:len(binary_payload)] & 0xFE) | np.array(list(binary_payload), dtype=np.uint8)
     stego_pixels = pixel_data.reshape((600, 800, 3))
+    
     return Image.fromarray(stego_pixels, 'RGB')
 
 def get_server_list():
-    status = f"Fetching server list from remote config..."
+    """Fetches and validates the server list from the remote JSON configuration."""
+    status = "Fetching server list from remote config..."
     yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, []
     try:
         response = requests.get(CREATOR_ENDPOINTS_JSON_URL, timeout=10)
         response.raise_for_status()
         all_entries = response.json()
         valid_endpoints = []
-        required_keys = ["name", "api_endpoint", "public_key"]
+        
         for entry in all_entries:
-            if all(key in entry for key in required_keys):
-                valid_endpoints.append(entry)
-            else:
-                logger.warning(f"Skipping invalid entry in configuration file: {entry}")
+            if not isinstance(entry, dict) or "name" not in entry or "public_key" not in entry:
+                logger.warning(f"Skipping invalid entry (missing name or public_key): {entry}")
+                continue
+
+            # **FIX: Auto-generate api_endpoint from link if it's missing**
+            if "api_endpoint" not in entry:
+                if "link" in entry:
+                    base_url = entry["link"].strip("/")
+                    entry["api_endpoint"] = f"{base_url}/api/predict/"
+                    logger.info(f"Auto-generating API endpoint for '{entry['name']}': {entry['api_endpoint']}")
+                else:
+                    logger.warning(f"Skipping invalid entry '{entry['name']}' (missing 'api_endpoint' and 'link'): {entry}")
+                    continue
+
+            valid_endpoints.append(entry)
+
         if not valid_endpoints:
             raise ValueError("No valid server configurations found in the remote file.")
+            
         endpoint_names = [e['name'] for e in valid_endpoints]
         status = f"✅ Success! Found {len(endpoint_names)} valid servers."
         yield gr.Dropdown(choices=endpoint_names, value=endpoint_names[0] if endpoint_names else None, label="Target Server"), status, valid_endpoints
     except Exception as e:
         status = f"❌ Error fetching or parsing configuration: {e}"
+        logger.error(status)
         yield gr.Dropdown(choices=[], value=None, label="Error fetching servers"), status, []
 
 def create_keylock_wrapper(service_name: str, secret_data: str, available_endpoints: list):
-    if not service_name: raise gr.Error("Please select a target server.")
+    """Wrapper function to handle UI for creating the encrypted image."""
+    if not service_name:
+        raise gr.Error("Please select a target server.")
     public_key = next((e['public_key'] for e in available_endpoints if e['name'] == service_name), None)
-    if not public_key: raise gr.Error(f"Could not find public key for '{service_name}'.")
+    if not public_key:
+        raise gr.Error(f"Could not find public key for '{service_name}'. Please refresh the server list.")
     try:
         created_image = create_encrypted_image(secret_data, public_key)
         return created_image, f"✅ Success! Image created for '{service_name}'."
     except Exception as e:
+        logger.error(f"Error creating image: {e}")
         return None, f"❌ Error: {e}"
 
 def send_keylock_wrapper(service_name: str, image: Image.Image, available_endpoints: list):
-    if not service_name: raise gr.Error("Please select a target server.")
-    if image is None: raise gr.Error("Please upload an image to send.")
+    """Wrapper function to handle UI for sending the image to the remote server."""
+    if not service_name:
+        raise gr.Error("Please select a target server.")
+    if image is None:
+        raise gr.Error("Please create or upload an image to send.")
+        
     api_endpoint = next((e.get('api_endpoint') for e in available_endpoints if e['name'] == service_name), None)
-    if not api_endpoint: raise gr.Error(f"Configuration Error: Could not find 'api_endpoint' for '{service_name}'.")
+    if not api_endpoint:
+        raise gr.Error(f"Configuration Error: Could not find 'api_endpoint' for '{service_name}'.")
+        
     status = f"Connecting to endpoint: {api_endpoint}"
     yield None, status
+    
     try:
         with io.BytesIO() as buffer:
             image.save(buffer, format="PNG")
             b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        
+        # Gradio API expects data in this format
         payload = {"data": [b64_string]}
         headers = {"Content-Type": "application/json"}
+        
         response = requests.post(api_endpoint, headers=headers, json=payload, timeout=45)
+        response.raise_for_status() # **IMPROVEMENT: Raise exception for HTTP errors**
+        
         response_json = response.json()
-        if response.status_code == 200:
-            if "data" in response_json:
-                yield response_json["data"][0], "✅ Success! Data decrypted by remote server."
-            else:
-                raise gr.Error(f"API returned an unexpected success format: {response_json}")
+        if "data" in response_json:
+            decrypted_data = response_json["data"][0]
+            # If the server returns a stringified JSON, parse it again
+            if isinstance(decrypted_data, str):
+                try:
+                    decrypted_data = json.loads(decrypted_data)
+                except json.JSONDecodeError:
+                    pass # Keep as string if not valid JSON
+            yield decrypted_data, "✅ Success! Data decrypted by remote server."
+        elif "error" in response_json:
+            raise gr.Error(f"API Error: {response_json['error']}")
         else:
-            raise gr.Error(f"API Error (Status {response.status_code}): {response_json.get('error', 'Unknown error')}")
+            raise gr.Error(f"API returned an unexpected response format: {response_json}")
+
     except Exception as e:
+        logger.error(f"Error calling server API: {e}")
         yield None, f"❌ Error calling server API: {e}"
 
 def refresh_and_update_all():
+    """Calls get_server_list and updates all relevant UI components."""
+    # The generator yields updates, we need to iterate to get the final one
     for dropdown_update, status_update, state_update in get_server_list():
         pass
     return dropdown_update, dropdown_update, status_update, state_update
 
+# --- Gradio UI Definition ---
 theme = gr.themes.Base(
     primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky, neutral_hue=gr.themes.colors.slate,
     font=(gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"),
@@ -186,6 +266,7 @@ with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo:
                         output_private_key = gr.Textbox(lines=10, label="Generated Private Key", interactive=False, show_copy_button=True)
                 gen_keys_button = gr.Button("⚙️ Generate New 2048-bit Key Pair", variant="secondary")
 
+    # --- Event Handlers ---
     gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key])
     refresh_button.click(fn=refresh_and_update_all, outputs=[creator_service_dropdown, send_service_dropdown, creator_status, endpoints_state])
     demo.load(fn=refresh_and_update_all, outputs=[creator_service_dropdown, send_service_dropdown, creator_status, endpoints_state])