Update app.py

app.py

@@ -12,84 +12,85 @@ logger = logging.getLogger(__name__)
 #  CONFIGURATION & DATA HANDLING
 # ==============================================================================
 
-ENDPOINTS = []
-ENDPOINT_NAMES = []
-DEFAULT_CONFIG_FILE = "endpoints.json"
+DEFAULT_CONFIG_FILE = "https://huggingface.co/spaces/Space-Share/bucket/raw/main/endpoints.json"
 
-def load_config_from_string(json_string: str) -> str:
-    """Parses a JSON string, updates global config, and returns a status message."""
-    global ENDPOINTS, ENDPOINT_NAMES
+def load_config_from_file(filepath: str) -> list:
+    """Loads and validates the JSON configuration from a file."""
     try:
-        data = json.loads(json_string)
+        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.")
-        
-        # Validate and extract data
-        temp_endpoints = []
-        temp_names = []
         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.")
-            temp_endpoints.append(service)
-            temp_names.append(service["name"])
+        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 an empty list on error to prevent crashing the app
+        return []
 
-        ENDPOINTS = temp_endpoints
-        ENDPOINT_NAMES = temp_names
-        
-        status_msg = f"✅ Success! Loaded {len(ENDPOINT_NAMES)} services from custom JSON."
+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 loading custom JSON: {e}"
+        status_msg = f"❌ Error saving configuration: {e}"
         logger.error(status_msg)
-        # On error, we don't change the existing valid configuration
         return status_msg
 
 # --- Load initial configuration from the file at startup ---
-try:
-    with open(DEFAULT_CONFIG_FILE, "r") as f:
-        initial_config = f.read()
-        INITIAL_STATUS = load_config_from_string(initial_config)
-except FileNotFoundError:
-    logger.warning(f"{DEFAULT_CONFIG_FILE} not found. App will start with an empty configuration.")
-    initial_config = "[]"
-    INITIAL_STATUS = f"Status: {DEFAULT_CONFIG_FILE} not found. Please provide a configuration."
-except Exception as e:
-    logger.error(f"Error loading initial config file: {e}")
-    initial_config = "[]"
-    INITIAL_STATUS = f"Status: Error loading {DEFAULT_CONFIG_FILE}: {e}"
+# The 'config_state' will hold our list of dictionaries
+initial_config = load_config_from_file(DEFAULT_CONFIG_FILE)
 
 # ==============================================================================
-#  MAIN APPLICATION LOGIC
+#  GRADIO UI HELPER FUNCTIONS
 # ==============================================================================
 
-def generate_image(selected_service_name: str, secret_data_str: str):
-    """
-    The main function called by the 'Create' button.
-    It finds the selected service's public key and calls the core encryption logic.
-    """
+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}
+    
+    # Check for duplicate names
+    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]
+    
+    # Update dropdown choices
+    updated_choices = [service['name'] for service in updated_config]
+    
+    # Clear the input forms and provide success feedback
+    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, now takes the config state as an input."""
     if not all([selected_service_name, secret_data_str]):
         raise gr.Error("A selected service and secret data are both required.")
 
     try:
-        # Find the public key for the selected service from the loaded global config
-        public_key = None
-        for service in ENDPOINTS:
-            if service.get("name") == selected_service_name:
-                public_key = service.get("public_key")
-                break
-        
+        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.")
 
-        # Call the core function to do the heavy lifting
         encrypted_image = core.create_encrypted_image(secret_data_str, public_key)
         
-        # Save to a temporary file for download
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_file:
             encrypted_image.save(tmp_file.name)
-            logger.info(f"Generated image saved to temporary file: {tmp_file.name}")
-            return encrypted_image, tmp_file.name, f"✅ Success! Image created for '{selected_service_name}'. You can now download it."
+            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)
@@ -100,22 +101,22 @@ def generate_image(selected_service_name: str, secret_data_str: str):
 # ==============================================================================
 
 with gr.Blocks(theme=gr.themes.Soft(), title="KeyLock Image Creator") as demo:
+    # This gr.State object holds the current configuration (list of dicts)
+    # It's initialized with the data loaded from endpoints.json
+    config_state = gr.State(initial_config)
+    
     gr.Markdown("# 🏭 KeyLock Image Creator")
     gr.Markdown("Create secure, encrypted PNG images for various decoder services.")
 
-    # This invisible textbox holds the current state of the dropdown choices
-    # This allows us to update the dropdown on the main tab from the config tab
-    dropdown_state = gr.State(ENDPOINT_NAMES)
-
-    with gr.Tabs():
+    with gr.Tabs() as tabs:
         with gr.TabItem("Create Image", id=0):
             with gr.Row():
                 with gr.Column(scale=2):
                     gr.Markdown("### 1. Select Decoder Service")
                     service_dropdown = gr.Dropdown(
                         label="Target Service",
-                        choices=ENDPOINT_NAMES,
-                        value=ENDPOINT_NAMES[0] if ENDPOINT_NAMES else None,
+                        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")
@@ -127,42 +128,65 @@ with gr.Blocks(theme=gr.themes.Soft(), title="KeyLock Image Creator") as demo:
                 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)
+                    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):
             gr.Markdown("## Decoder Service Configuration")
-            gr.Markdown("You can override the default `endpoints.json` by pasting a new JSON configuration below.")
-            json_config = gr.Textbox(
-                lines=15,
-                label="JSON Configuration",
-                value=initial_config,
-                info="Each service must be an object in a list with 'name', 'link', and 'public_key'."
-            )
-            update_config_button = gr.Button("Apply New Configuration")
-            config_status_output = gr.Textbox(label="Configuration Status", value=INITIAL_STATUS, interactive=False)
+            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="Loaded from endpoints.json")
+                
+                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")
+            config_status_output = gr.Textbox(label="Configuration Status", interactive=False)
 
     # --- Interactivity ---
-    # When the "Apply" button on the config tab is clicked...
-    update_config_button.click(
-        fn=load_config_from_string,
-        inputs=[json_config],
-        outputs=[config_status_output]
-    ).then(
-        # ...then, update the dropdown on the main tab with the new choices.
-        # We use a lambda to get the latest global state.
-        lambda: gr.Dropdown(choices=ENDPOINT_NAMES, value=ENDPOINT_NAMES[0] if ENDPOINT_NAMES else None),
-        inputs=[],
-        outputs=[service_dropdown]
-    )
-
-    # When the "Create" button on the main tab is clicked...
+    
+    # Main image creation logic
     create_button.click(
         fn=generate_image,
-        inputs=[service_dropdown, secret_data],
+        inputs=[service_dropdown, secret_data, config_state],
         outputs=[image_output, download_output, status_output]
     )
 
+    # Logic for the "Manage Configuration" tab
+    add_service_button.click(
+        fn=add_new_service,
+        inputs=[config_state, new_service_name, new_service_link, new_service_key],
+        # Outputs: update the state, the dropdown, the status, and clear the form fields
+        outputs=[config_state, service_dropdown, config_status_output, new_service_name, new_service_link, new_service_key]
+    ).then(
+        # After updating the state, also update the JSON display to reflect the change
+        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]
+    )
+
+    # When the configuration state changes (e.g., after adding a service),
+    # update the JSON display on the config tab. This keeps the UI in sync.
+    config_state.change(
+        fn=lambda state: state,
+        inputs=[config_state],
+        outputs=[config_display]
+    )
+
 if __name__ == "__main__":
     demo.launch()