import gradio as gr from gradio_client import Client from PIL import Image import base64 import io import json import logging import requests # Use requests to fetch the remote JSON file from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa # --- Configure Logging --- logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # ============================================================================== # CONFIGURATION: IDs AND URLs OF THE REMOTE SERVICES # ============================================================================== CREATOR_SPACE_ID = "broadfield-dev/KeyLock-Auth-Creator" SERVER_SPACE_ID = "broadfield-dev/KeyLock-Auth-Server" # URL to the raw JSON file containing the list of public keys and services. # This makes the Creator's configuration publicly readable. CREATOR_ENDPOINTS_JSON_URL = "https://huggingface.co/spaces/Space-Share/bucket/raw/main/endpoints.json" # Construct URLs for linking in documentation BASE_HF_URL = "https://huggingface.co/spaces/" CREATOR_URL = f"{BASE_HF_URL}{CREATOR_SPACE_ID}" SERVER_URL = f"{BASE_HF_URL}{SERVER_SPACE_ID}" CREATOR_APP_PY_URL = f"{CREATOR_URL}/blob/main/app.py" SERVER_APP_PY_URL = f"{SERVER_URL}/blob/main/app.py" # ============================================================================== # API CALL WRAPPER FUNCTIONS # ============================================================================== def get_creator_endpoints(): """Fetches the list of supported endpoints by making an HTTP request to the Creator's JSON file.""" status = f"Fetching endpoint list from {CREATOR_ENDPOINTS_JSON_URL}..." yield gr.Dropdown(choices=[], value=None, label="⏳ Fetching..."), status, [] # Initial state try: response = requests.get(CREATOR_ENDPOINTS_JSON_URL, timeout=10) response.raise_for_status() # Raise an exception for bad status codes endpoints = response.json() endpoint_names = [e['name'] for e in endpoints] status = f"✅ Success! Found {len(endpoint_names)} endpoints." # Return the full list to the state, and the updated dropdown yield gr.Dropdown(choices=endpoint_names, value=endpoint_names[0] if endpoint_names else None, label="Target Service"), status, endpoints except Exception as e: logger.error(f"Failed to get endpoints from creator's JSON file: {e}", exc_info=True) status = f"❌ Error: Could not fetch configuration. Check the URL and if the 'endpoints.json' file is public. Details: {e}" yield gr.Dropdown(choices=[], value=None, label="Error fetching services"), status, [] def create_image_via_api(service_name: str, secret_data: str, available_endpoints: list): """Calls the Creator Space API to generate an encrypted image for a selected service.""" if not all([service_name, secret_data]): raise gr.Error("Please select a service and provide secret data.") status = f"Looking up public key for '{service_name}'..." yield None, None, status 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}' in the fetched configuration.") status = f"Connecting to Creator: {CREATOR_SPACE_ID}..." yield None, None, status try: client = Client(src=CREATOR_SPACE_ID) temp_filepath = client.predict(secret_data, public_key, api_name="/create_image") if not temp_filepath: raise gr.Error("Creator API did not return an image.") created_image = Image.open(temp_filepath) status = f"✅ Success! Image created for '{service_name}'." yield created_image, temp_filepath, status except Exception as e: logger.error(f"Creator API call failed: {e}", exc_info=True) yield None, None, f"❌ Error calling Creator API: {e}" def decrypt_image_via_api(image: Image.Image): """Calls the Server Space API to decrypt an image.""" if image is None: raise gr.Error("Please upload an image to decrypt.") status = f"Connecting to Server: {SERVER_SPACE_ID}..." yield None, status try: client = Client(src=SERVER_SPACE_ID) with io.BytesIO() as buffer: image.save(buffer, format="PNG") b64_string = base64.b64encode(buffer.getvalue()).decode("utf-8") status = f"Calling API on {SERVER_SPACE_ID}..." yield None, status decrypted_json = client.predict(b64_string, api_name="/keylock-auth-decoder") status = "✅ Success! Data decrypted by the Server." yield decrypted_json, status except Exception as e: logger.error(f"Server API call failed: {e}", exc_info=True) yield None, f"❌ Error calling Server API: {e}" def generate_rsa_keys(): """Generates a new RSA 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 # ============================================================================== # GRADIO DASHBOARD INTERFACE # ============================================================================== 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"), ).set( body_background_fill="#F1F5F9", panel_background_fill="white", block_background_fill="white", block_border_width="1px", block_shadow="*shadow_drop_lg", button_primary_background_fill="*primary_600", button_primary_background_fill_hover="*primary_700", ) with gr.Blocks(theme=theme, title="KeyLock Operations Dashboard") as demo: endpoints_state = gr.State([]) gr.Markdown("# 🔑 KeyLock Operations Dashboard") gr.Markdown("A centralized dashboard to manage and demonstrate the entire KeyLock ecosystem, powered by live API calls to dedicated services.") with gr.Tabs() as tabs: with gr.TabItem("① Generate Keys", id=0): gr.Markdown("## RSA Key Pair Generator") gr.Markdown("Create a new public/private key pair. The public key can be added to a service's configuration to allow it to be a target for the Auth Creator.") with gr.Row(variant="panel"): with gr.Group(): gen_keys_button = gr.Button("🔑 Generate New 2048-bit Key Pair", variant="secondary") with gr.Row(): output_public_key = gr.Textbox(lines=10, label="Generated Public Key (Share This)", interactive=False, show_copy_button=True) output_private_key = gr.Textbox(lines=10, label="Generated Private Key (Keep Secret!)", interactive=False, show_copy_button=True) with gr.TabItem("② Auth Creator", id=1): gr.Markdown("## Create an Encrypted Authentication Image") gr.Markdown(f"This tool calls the **[{CREATOR_SPACE_ID}]({CREATOR_URL})** service to encrypt data for a chosen target. The list of targets is fetched live from the Creator's configuration.") with gr.Row(variant="panel"): with gr.Column(scale=2): with gr.Row(): creator_service_dropdown = gr.Dropdown(label="Target Service", interactive=True, info="Select the API server you want to encrypt data for.") refresh_button = gr.Button("🔄 Refresh List", scale=0, size="sm") creator_secret_input = gr.Textbox(lines=8, label="Secret Data to Encrypt", placeholder="API_KEY: sk-123...\nUSER: demo-user") creator_button = gr.Button("✨ Create Auth Image via API", variant="primary") with gr.Column(scale=1): creator_status = gr.Textbox(label="Status", interactive=False, lines=2) creator_image_output = gr.Image(label="Image from Creator Service", type="pil", show_download_button=True) with gr.TabItem("③ Client / Decoder", id=2): gr.Markdown("## Decrypt an Authentication Image") gr.Markdown(f"This tool acts as a client, calling the **[{SERVER_SPACE_ID}]({SERVER_URL})** service to decrypt an image using its securely stored private key.") with gr.Row(variant="panel"): with gr.Column(scale=1): client_image_input = gr.Image(type="pil", label="Upload Encrypted Auth Image", sources=["upload", "clipboard"]) client_button = gr.Button("🔓 Decrypt Image via Server API", variant="primary") with gr.Column(scale=1): client_status = gr.Textbox(label="Status", interactive=False, lines=2) client_json_output = gr.JSON(label="Decrypted Data from Server") with gr.TabItem("ℹ️ Service Information", id=3): gr.Markdown("## Ecosystem Architecture") gr.Markdown("This dashboard coordinates separate Hugging Face Spaces to demonstrate a secure, decoupled workflow. Each service has a specific role.") with gr.Row(): with gr.Column(): gr.Markdown(f"### 🏭 Auth Creator Service\n- **Space:** [{CREATOR_SPACE_ID}]({CREATOR_URL})\n- **Role:** Provides an API to encrypt data for various targets defined in its `endpoints.json` file.\n- **Source Code:** [app.py]({CREATOR_APP_PY_URL})") with gr.Column(): gr.Markdown(f"### 📡 Decoder Server\n- **Space:** [{SERVER_SPACE_ID}]({SERVER_URL})\n- **Role:** The trusted authority. It holds a secret private key and provides a secure API to decrypt images.\n- **Source Code:** [app.py]({SERVER_APP_PY_URL})") # --- Wire up the component logic --- gen_keys_button.click(fn=generate_rsa_keys, inputs=None, outputs=[output_private_key, output_public_key]) # Event handler for loading the page or refreshing the endpoint list def refresh_endpoints(): # This is a generator function, so we need to iterate to get the last value. *_, last_yield = get_creator_endpoints() return last_yield refresh_button.click(fn=refresh_endpoints, outputs=[creator_service_dropdown, creator_status, endpoints_state]) demo.load(fn=refresh_endpoints, outputs=[creator_service_dropdown, creator_status, endpoints_state]) # Logic for the Creator and Client tabs creator_button.click(fn=create_image_via_api, inputs=[creator_service_dropdown, creator_secret_input, endpoints_state], outputs=[creator_image_output, gr.File(visible=False), creator_status]) client_button.click(fn=decrypt_image_via_api, inputs=[client_image_input], outputs=[client_json_output, client_status]) if __name__ == "__main__": demo.launch()