Update app.py

app.py

@@ -10,7 +10,7 @@ import shodan
 import time
 from datasets import load_dataset, Dataset
 from huggingface_hub import HfApi, login
-from huggingface_hub import secrets
+import os
 from typing import List, Dict, Any, Optional, Tuple
 import pandas as pd
 
@@ -48,7 +48,7 @@ def get_dataset() -> Dataset:
     """Get the dataset with Ollama instances and models."""
     try:
         # Try to load the existing dataset
-        hf_token = secrets.get_secret("HF_TOKEN")
+        hf_token = os.getenv("HF_TOKEN")
         dataset = load_dataset(DATASET_NAME, use_auth_token=hf_token)
         logger.info(f"Loaded existing dataset: {DATASET_NAME}")
         return dataset["train"]
@@ -70,7 +70,7 @@ def get_dataset() -> Dataset:
 def push_dataset(dataset: Dataset) -> None:
     """Push the dataset to the Hub."""
     try:
-        hf_token = secrets.get_secret("HF_TOKEN")
+        hf_token = os.getenv("HF_TOKEN")
         dataset.push_to_hub(DATASET_NAME, token=hf_token)
         logger.info(f"Pushed dataset to Hub: {DATASET_NAME}")
     except Exception as e:
@@ -203,8 +203,13 @@ async def scan_ollama_instances(progress=None) -> Tuple[str, Dataset]:
     
     # Get secrets
     try:
-        shodan_api_key = secrets.get_secret("SHODAN_API_KEY")
-        shodan_query = secrets.get_secret("SHODAN_QUERY")
+        shodan_api_key = os.getenv("SHODAN_API_KEY")
+        shodan_query = os.getenv("SHODAN_QUERY")
+        
+        if not shodan_api_key:
+            raise ValueError("SHODAN_API_KEY environment variable is not set")
+        if not shodan_query:
+            raise ValueError("SHODAN_QUERY environment variable is not set")
     except Exception as e:
         logger.error(f"Failed to get secrets: {e}")
         return "Error: Failed to retrieve secrets", None
@@ -370,10 +375,195 @@ def get_unique_values(dataset: Dataset) -> Tuple[List[str], List[str]]:
 def login_submit(password: str) -> Tuple[bool, str, str]:
     """Handle admin login."""
     try:
-        stored_password = secrets.get_secret("ADMIN_PASSWORD")
+        stored_password = os.getenv("ADMIN_PASSWORD")
+        
+        if not stored_password:
+            logger.error("ADMIN_PASSWORD environment variable is not set")
+            return False, "error", "Admin password is not configured. Please contact the administrator."
         
         # Check if stored password is already hashed
-        if stored_password.startswith('$2b$'):
+        if stored_password.startswith('$2b
+
+def search_models(
+    family: str, 
+    parameter_size: str, 
+    name_search: str, 
+    is_admin: bool,
+    dataset: Dataset
+) -> Tuple[pd.DataFrame, List[Dict[str, Any]]]:
+    """Search and filter models in the dataset."""
+    df, details = filter_and_search_models(
+        dataset,
+        family=None if family == "All" else family,
+        parameter_size=None if parameter_size == "All" else parameter_size,
+        name_search=name_search,
+        is_admin=is_admin
+    )
+    return df, details
+
+def show_model_details(evt: gr.SelectData, models: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """Show details for a selected model."""
+    if not models or evt.index[0] >= len(models):
+        return {}
+    
+    model = models[evt.index[0]]
+    return model
+
+# Main Gradio application
+def create_app():
+    with gr.Blocks(title="Ollama Scanner", theme=gr.themes.Soft()) as app:
+        # State variables
+        admin_logged_in = gr.State(False)
+        dataset_state = gr.State(get_dataset())
+        model_details_state = gr.State([])
+        
+        # Header
+        gr.Markdown("# 🔍 Ollama Scanner")
+        gr.Markdown("Browse publicly accessible Ollama instances and their models")
+        
+        # Login tab
+        with gr.Tab("Admin Login") as login_tab:
+            with gr.Group():
+                gr.Markdown("### Admin Login")
+                gr.Markdown("Enter the admin password to access administrative features.")
+                
+                admin_password = gr.Textbox(
+                    type="password", 
+                    label="Admin Password", 
+                    placeholder="Enter admin password"
+                )
+                login_btn = gr.Button("Login", variant="primary")
+                login_status = gr.Markdown("")
+        
+        # Browse Models tab
+        with gr.Tab("Browse Models") as browse_tab:
+            with gr.Row():
+                # Filters column
+                with gr.Column(scale=1):
+                    gr.Markdown("### Filters")
+                    
+                    family_dropdown = gr.Dropdown(
+                        choices=["All"], 
+                        value="All",
+                        label="Model Family"
+                    )
+                    
+                    parameter_size_dropdown = gr.Dropdown(
+                        choices=["All"], 
+                        value="All",
+                        label="Parameter Size"
+                    )
+                    
+                    name_search = gr.Textbox(
+                        label="Search by Name",
+                        placeholder="Search model names..."
+                    )
+                    
+                    search_btn = gr.Button("Search", variant="primary")
+                    
+                    stats_box = gr.Markdown("Loading stats...")
+                
+                # Results column
+                with gr.Column(scale=2):
+                    gr.Markdown("### Results")
+                    
+                    results_table = gr.DataFrame(
+                        label="Models",
+                        interactive=False
+                    )
+                    
+                    model_json_display = gr.JSON(
+                        label="Model Details",
+                        visible=True
+                    )
+        
+        # Shodan Scan tab (admin only)
+        with gr.Tab("Shodan Scan", visible=False) as scan_tab:
+            with gr.Group():
+                gr.Markdown("### Shodan Scan")
+                gr.Markdown("Scan for publicly accessible Ollama instances using Shodan.")
+                
+                scan_btn = gr.Button("Start Scan", variant="primary")
+                scan_progress = gr.Textbox(
+                    label="Scan Status",
+                    placeholder="Click 'Start Scan' to begin scanning...",
+                    interactive=False
+                )
+        
+        # Login logic
+        def handle_login(password):
+            is_admin, status, message = login_submit(password)
+            if is_admin:
+                return is_admin, message, gr.update(visible=True)
+            else:
+                return is_admin, message, gr.update(visible=False)
+        
+        login_btn.click(
+            handle_login,
+            inputs=[admin_password],
+            outputs=[admin_logged_in, login_status, scan_tab]
+        )
+        
+        # Search logic
+        def update_stats(dataset):
+            total_instances = len(dataset)
+            models_count = sum(len(instance.get('models', [])) for instance in dataset)
+            
+            families, parameter_sizes = get_unique_values(dataset)
+            family_count = len(families) - 1  # Subtract "All"
+            
+            return f"**Stats:** {total_instances} instances, {models_count} models, {family_count} families"
+        
+        search_btn.click(
+            search_models,
+            inputs=[family_dropdown, parameter_size_dropdown, name_search, admin_logged_in, dataset_state],
+            outputs=[results_table, model_details_state]
+        )
+        
+        # Model selection logic
+        results_table.select(
+            show_model_details,
+            inputs=[model_details_state],
+            outputs=[model_json_display]
+        )
+        
+        # Scan logic
+        async def run_scan():
+            result, updated_dataset = await scan_ollama_instances()
+            if updated_dataset is not None:
+                return result, updated_dataset, *get_unique_values(updated_dataset), update_stats(updated_dataset)
+            else:
+                return result, None, [], [], ""
+        
+        scan_btn.click(
+            run_scan,
+            inputs=[],
+            outputs=[scan_progress, dataset_state, family_dropdown, parameter_size_dropdown, stats_box]
+        )
+        
+        # Initial data load
+        def init_ui(dataset):
+            families, parameter_sizes = get_unique_values(dataset)
+            stats = update_stats(dataset)
+            
+            # Run initial search
+            df, details = search_models("All", "All", "", False, dataset)
+            
+            return families, parameter_sizes, stats, df, details
+        
+        app.load(
+            init_ui,
+            inputs=[dataset_state],
+            outputs=[family_dropdown, parameter_size_dropdown, stats_box, results_table, model_details_state]
+        )
+    
+    return app
+
+# Start the application
+if __name__ == "__main__":
+    app = create_app()
+    app.launch()
+):
             is_valid = verify_password(password, stored_password)
         else:
             # Compare directly for first-time setup
@@ -383,6 +573,8 @@ def login_submit(password: str) -> Tuple[bool, str, str]:
             if is_valid:
                 hashed_password = hash_password(stored_password)
                 logger.info("Hashed admin password for future use")
+                # Note: We can't store the hashed password back to environment variables
+                # in a Hugging Face Space environment. This would require a different approach.
         
         if is_valid:
             return True, "success", "Login successful! You now have admin access."