import gradio as gr import json import logging import tempfile import core # Import our core module # --- Configure Logging --- logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # ============================================================================== # CONFIGURATION & DATA HANDLING # ============================================================================== DEFAULT_CONFIG_FILE = "./endpoints.json" def load_config_from_file(filepath: str) -> list: """Loads and validates the JSON configuration from a file.""" try: with open(filepath, "r") as f: data = json.load(f) if not isinstance(data, list): raise TypeError("JSON root must be a list of service objects.") for service in data: if not all(k in service for k in ("name", "link", "public_key")): raise ValueError("Each service object must contain 'name', 'link', and 'public_key' keys.") logger.info(f"Successfully loaded {len(data)} services from {filepath}.") return data except FileNotFoundError: logger.warning(f"{filepath} not found. Returning empty list.") return [] except Exception as e: logger.error(f"Error loading config file {filepath}: {e}") return [] def save_config_to_file(filepath: str, config_data: list) -> str: """Saves the current configuration data to a JSON file.""" try: with open(filepath, "w") as f: json.dump(config_data, f, indent=2) status_msg = f"✅ Success! Configuration saved to {filepath}." logger.info(status_msg) return status_msg except Exception as e: status_msg = f"❌ Error saving configuration: {e}" logger.error(status_msg) return status_msg initial_config = load_config_from_file(DEFAULT_CONFIG_FILE) # ============================================================================== # GRADIO UI HELPER FUNCTIONS # ============================================================================== def add_new_service(current_config: list, name: str, link: str, public_key: str): """Adds a new service to the current configuration state.""" if not all([name, link, public_key]): raise gr.Error("All fields (Name, Link, Public Key) are required to add a new service.") new_service = {"name": name, "link": link, "public_key": public_key} if any(service['name'] == name for service in current_config): raise gr.Error(f"A service with the name '{name}' already exists. Please use a unique name.") updated_config = current_config + [new_service] updated_choices = [service['name'] for service in updated_config] status_update = f"✅ '{name}' added. You can now use it in the 'Create Image' tab. Click 'Save to File' to make it permanent." return updated_config, gr.Dropdown(choices=updated_choices), status_update, "", "", "" def generate_image(selected_service_name: str, secret_data_str: str, current_config: list): """The main image creation function for the UI.""" if not all([selected_service_name, secret_data_str]): raise gr.Error("A selected service and secret data are both required.") try: public_key = next((s['public_key'] for s in current_config if s.get('name') == selected_service_name), None) if not public_key: raise gr.Error(f"Could not find the service '{selected_service_name}'. Please check the configuration.") # This call is correct and should not cause an error encrypted_image = core.create_encrypted_image(secret_data_str, public_key) with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file: encrypted_image.save(tmp_file.name) logger.info(f"Generated image saved for '{selected_service_name}'.") return encrypted_image, tmp_file.name, f"✅ Success! Image created for '{selected_service_name}'." except Exception as e: logger.error(f"Image creation failed: {e}", exc_info=True) return None, None, f"❌ Error: {e}" def get_endpoints_json() -> str: """Reads and returns the content of endpoints.json as a string.""" try: with open(DEFAULT_CONFIG_FILE, "r") as f: return f.read() except Exception as e: logger.error(f"Could not read {DEFAULT_CONFIG_FILE}: {e}") raise gr.Error(f"Server could not read its {DEFAULT_CONFIG_FILE} configuration file.") # ============================================================================== # GRADIO INTERFACE # ============================================================================== with gr.Blocks(theme=gr.themes.Soft(), title="KeyLock Image Creator") as demo: config_state = gr.State(initial_config) gr.Markdown("# 🏭 KeyLock Image Creator") gr.Markdown("Create secure, encrypted PNG images for various decoder services.") with gr.Tabs() as tabs: with gr.TabItem("Create Image", id=0): # ... (Your existing UI code is fine, no changes needed here) ... with gr.Row(): with gr.Column(scale=2): gr.Markdown("### 1. Select Decoder Service") service_dropdown = gr.Dropdown( label="Target Service", choices=[s['name'] for s in initial_config], value=[s['name'] for s in initial_config][0] if initial_config else None, interactive=True ) gr.Markdown("### 2. Enter Your Secret Data") secret_data = gr.Textbox( lines=7, label="Secret Data (Key-Value Pairs)", placeholder="USERNAME: user@example.com\nAPI_KEY: sk-12345..." ) with gr.Column(scale=1): gr.Markdown("### 3. Generate") create_button = gr.Button("Create Encrypted Image", variant="primary") status_output = gr.Textbox(label="Status", interactive=False, lines=2) image_output = gr.Image(label="Generated Image", type="pil") download_output = gr.File(label="Download Image") with gr.TabItem("Manage Configuration", id=1): # ... (Your existing config management code is fine, no changes needed here) ... gr.Markdown("## Decoder Service Configuration") gr.Markdown("View the current configuration, add new services, and save your changes.") with gr.Row(): with gr.Column(scale=2): gr.Markdown("### Current Configuration") config_display = gr.JSON(value=initial_config, label=f"Loaded from {DEFAULT_CONFIG_FILE}") with gr.Column(scale=1): gr.Markdown("### Add a New Service") new_service_name = gr.Textbox(label="Service Name", placeholder="e.g., My Production API") new_service_link = gr.Textbox(label="Service Link", placeholder="https://huggingface.co/...") new_service_key = gr.Textbox(label="Public Key (PEM Format)", lines=5, placeholder="-----BEGIN PUBLIC KEY-----...") add_service_button = gr.Button("Add Service to Current Session") gr.Markdown("---") save_config_button = gr.Button("Save Current Configuration to File", variant="secondary", visible=False) config_status_output = gr.Textbox(label="Configuration Status", interactive=False) # FIX: Define API endpoints clearly instead of using a hidden row with gr.TabItem("API", id=2): gr.Markdown("## API Endpoints") gr.Markdown("Use these endpoints for programmatic access.") # API for creating an image gr.Interface( # By passing core.create_encrypted_image directly, we ensure it runs in the correct scope fn=core.create_encrypted_image, inputs=[ gr.Textbox(label="Secret Data", info="Key-value pairs, one per line."), gr.Textbox(label="Public Key PEM", info="The full PEM-formatted public key.") ], outputs=gr.Image(type="pil", label="Encrypted Image"), title="Image Creation API", description="Creates an encrypted image from secret data and a public key. Use via the `/api/create_image/` route.", api_name="create_image" ) # API to get the endpoint list gr.Interface( fn=get_endpoints_json, inputs=[], outputs=gr.Textbox(label="Endpoint Configuration"), title="Get Endpoints API", description=f"Retrieves the contents of `{DEFAULT_CONFIG_FILE}`. Use via the `/api/get_endpoints/` route.", api_name="get_endpoints" ) # --- Interactivity --- create_button.click( fn=generate_image, inputs=[service_dropdown, secret_data, config_state], outputs=[image_output, download_output, status_output] ) add_service_button.click( fn=add_new_service, inputs=[config_state, new_service_name, new_service_link, new_service_key], outputs=[config_state, service_dropdown, config_status_output, new_service_name, new_service_link, new_service_key] ).then( lambda state: state, inputs=[config_state], outputs=[config_display] ) save_config_button.click( fn=save_config_to_file, inputs=[gr.State(DEFAULT_CONFIG_FILE), config_state], outputs=[config_status_output] ) config_state.change( fn=lambda state: state, inputs=[config_state], outputs=[config_display] ) if __name__ == "__main__": demo.launch()