Update app.py

app.py

@@ -1,564 +1,532 @@
 import os
 import logging
-import requests
+import asyncio
 import bcrypt
-import gradio as gr
-from huggingface_hub import HfApi, login
-from datasets import load_dataset, Dataset, Features, Value, Sequence
-from typing import Dict, List, Optional, Any
-import time
-from concurrent.futures import ThreadPoolExecutor, as_completed
+import requests
 import shodan
-import html
+import gradio as gr
+from typing import List, Dict, Any, Tuple, Optional
+from datasets import load_dataset, Dataset
+from huggingface_hub import HfApi, create_repo
 
-# Configure logging
+# Setup logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
 
-# Function to access secrets
-def get_secret(name: str) -> str:
-    """Get a secret from Hugging Face Space secrets."""
-    try:
-        from huggingface_hub import secrets
-        return secrets.get_secret(name)
-    except Exception as e:
-        logger.warning(f"Error accessing secret {name}: {e}")
-        return ""
-
-# Sanitize inputs
-def sanitize_input(text: str) -> str:
-    """Sanitize user input to prevent injection attacks."""
-    if not text:
-        return ""
-    return html.escape(text)
+# Environment variable validation
+required_env_vars = ["SHODAN_API_KEY", "ADMIN_PASSWORD", "SHODAN_QUERY"]
+for var in required_env_vars:
+    if not os.environ.get(var):
+        logger.warning(f"Environment variable {var} is not set")
 
 # Dataset functions
-def get_or_create_dataset(repo_id: str, token: Optional[str] = None) -> Dataset:
+def get_or_create_dataset(dataset_name: str = "latterworks/llama_checker_results") -> Optional[Dataset]:
     """
-    Get or create a dataset with the specified schema.
+    Load the dataset or create it if it doesn't exist.
     
     Args:
-        repo_id: The Hugging Face dataset repository ID
-        token: Authentication token for private datasets
+        dataset_name: The name of the dataset on Hugging Face Hub
         
     Returns:
-        The dataset object
+        The dataset or None if there was an error
     """
     try:
         # Try to load the dataset
-        dataset = load_dataset(repo_id, token=token)
-        return dataset["train"]
-    except Exception as e:
-        logger.info(f"Dataset {repo_id} not found or error loading it: {e}")
-        
-        # Create a new dataset with the required schema
-        features = Features({
-            "ip": Value("string"),
-            "port": Value("int32"),
-            "country": Value("string", default_value=""),
-            "region": Value("string", default_value=""),
-            "org": Value("string", default_value=""),
-            "models": Sequence({
-                "name": Value("string"),
-                "family": Value("string", default_value=""),
-                "parameter_size": Value("string", default_value=""),
-                "quantization_level": Value("string", default_value=""),
-                "digest": Value("string", default_value=""),
-                "modified_at": Value("string", default_value=""),
-                "size": Value("int64", default_value=0)
-            }, default_value=[])
-        })
-        
-        # Create empty dataset
-        empty_dataset = Dataset.from_dict({
-            "ip": [],
-            "port": [],
-            "country": [],
-            "region": [],
-            "org": [],
-            "models": []
-        }, features=features)
+        token = os.environ.get("HF_TOKEN")
         
-        # Push to hub
-        empty_dataset.push_to_hub(repo_id, token=token)
-        
-        return empty_dataset
+        try:
+            dataset_dict = load_dataset(dataset_name, token=token)
+            if "train" in dataset_dict:
+                return dataset_dict["train"]
+            else:
+                # If no "train" split, try to use the first available split
+                first_split = next(iter(dataset_dict))
+                return dataset_dict[first_split]
+        except Exception as e:
+            logger.error(f"Failed to load dataset {dataset_name}: {e}")
+            
+            # Create the dataset
+            try:
+                # Create the repository
+                hf_api = HfApi(token=token)
+                create_repo(repo_id=dataset_name, repo_type="dataset", token=token)
+                
+                # Create empty dataset with the correct schema
+                empty_dataset = Dataset.from_dict({
+                    "ip": [],
+                    "port": [],
+                    "country": [],
+                    "region": [],
+                    "org": [],
+                    "models": []
+                })
+                
+                # Push to Hub
+                empty_dataset.push_to_hub(dataset_name, token=token)
+                
+                return empty_dataset
+            except Exception as create_e:
+                logger.error(f"Failed to create dataset: {create_e}")
+                return None
+    except Exception as e:
+        logger.exception(f"Unexpected error in get_or_create_dataset: {e}")
+        return None
 
-def update_dataset_entry(dataset: Dataset, entry: Dict[str, Any]) -> Dataset:
+def update_dataset(dataset: Dataset, new_entries: List[Dict[str, Any]]) -> Optional[Dataset]:
     """
-    Add or update an entry in the dataset.
+    Update the dataset with new entries from Shodan scan.
     
     Args:
         dataset: The dataset to update
-        entry: The entry to add or update
+        new_entries: List of new entries to add or update in the dataset
         
     Returns:
-        The updated dataset
+        The updated dataset or None if there was an error
     """
-    # Check if entry already exists by IP and port
-    ip = entry["ip"]
-    port = entry["port"]
-    
-    # Find the index if it exists
-    existing_idx = None
-    for idx, item in enumerate(dataset):
-        if item["ip"] == ip and item["port"] == port:
-            existing_idx = idx
-            break
-    
-    if existing_idx is not None:
-        # Update existing entry
-        dataset = dataset.select(list(range(len(dataset))))
-        new_examples = list(dataset)
-        new_examples[existing_idx] = entry
-        return Dataset.from_dict({k: [ex[k] for ex in new_examples] for k in dataset.column_names})
-    else:
-        # Add new entry
-        new_dataset = dataset.add_item(entry)
-        return new_dataset
-
-def push_dataset_to_hub(dataset: Dataset, repo_id: str, token: Optional[str] = None):
-    """
-    Push dataset to the Hugging Face Hub.
+    if dataset is None:
+        logger.error("Cannot update None dataset")
+        return None
     
-    Args:
-        dataset: The dataset to push
-        repo_id: The repository ID
-        token: Authentication token
-    """
-    try:
-        dataset.push_to_hub(repo_id, token=token)
-        logger.info(f"Successfully pushed dataset to {repo_id}")
-    except Exception as e:
-        logger.error(f"Error pushing dataset to hub: {e}")
-
-# Shodan functions
-def scan_with_shodan(shodan_api_key: str, query: str, max_results: int = 1000) -> List[Dict[str, Any]]:
-    """
-    Scan with Shodan API for Ollama instances.
+    # Convert dataset to dictionaries for easier manipulation
+    dataset_dict = dataset.to_dict()
     
-    Args:
-        shodan_api_key: Shodan API key
-        query: Shodan search query
-        max_results: Maximum number of results to return
+    # Add new entries
+    for entry in new_entries:
+        ip = entry.get("ip")
+        port = entry.get("port")
         
-    Returns:
-        List of discovered instances
-    """
-    if not shodan_api_key:
-        logger.error("No Shodan API key provided")
-        return []
-        
-    try:
-        api = shodan.Shodan(shodan_api_key)
-        results = []
-        
-        # Get the number of total results
-        count_result = api.count(query)
-        total_results = count_result['total']
-        logger.info(f"Found {total_results} results for query: {query}")
-        
-        # Limit to max_results
-        pages = min(total_results, max_results) // 100
-        if min(total_results, max_results) % 100 > 0:
-            pages += 1
-        
-        for page in range(1, pages + 1):
-            try:
-                result_page = api.search(query, page=page)
-                for match in result_page['matches']:
-                    instance = {
-                        "ip": match.get("ip_str", ""),
-                        "port": match.get("port", 11434),  # Default Ollama port
-                        "country": match.get("location", {}).get("country_name", ""),
-                        "region": match.get("location", {}).get("region_name", ""),
-                        "org": match.get("org", ""),
-                        "models": []  # Will be populated later
-                    }
-                    results.append(instance)
-                logger.info(f"Processed page {page}/{pages}")
-            except shodan.APIError as e:
-                logger.error(f"Shodan API error on page {page}: {e}")
+        # Check if this IP:port combination already exists
+        found = False
+        for i, (existing_ip, existing_port) in enumerate(zip(dataset_dict["ip"], dataset_dict["port"])):
+            if existing_ip == ip and existing_port == port:
+                # Update the entry
+                dataset_dict["country"][i] = entry.get("country", dataset_dict["country"][i])
+                dataset_dict["region"][i] = entry.get("region", dataset_dict["region"][i])
+                dataset_dict["org"][i] = entry.get("org", dataset_dict["org"][i])
+                dataset_dict["models"][i] = entry.get("models", dataset_dict["models"][i])
+                found = True
                 break
         
-        return results
-    except shodan.APIError as e:
-        logger.error(f"Shodan API error: {e}")
-        return []
+        if not found:
+            # Add as a new entry
+            dataset_dict["ip"].append(entry.get("ip", ""))
+            dataset_dict["port"].append(entry.get("port", 0))
+            dataset_dict["country"].append(entry.get("country", ""))
+            dataset_dict["region"].append(entry.get("region", ""))
+            dataset_dict["org"].append(entry.get("org", ""))
+            dataset_dict["models"].append(entry.get("models", []))
+    
+    # Convert back to Dataset
+    updated_dataset = Dataset.from_dict(dataset_dict)
+    
+    # Push to Hub
+    token = os.environ.get("HF_TOKEN")
+    updated_dataset.push_to_hub("latterworks/llama_checker_results", token=token)
+    
+    return updated_dataset
 
 # Ollama endpoint checking
-def check_ollama_endpoint(instance: Dict[str, Any], timeout: int = 5) -> Dict[str, Any]:
+async def check_ollama_endpoint(ip: str, port: int) -> Dict[str, Any]:
     """
-    Check an Ollama endpoint and extract model information.
+    Check a single Ollama endpoint and retrieve model information.
     
     Args:
-        instance: Instance information (ip, port, etc.)
-        timeout: Request timeout in seconds
+        ip: The IP address of the Ollama instance
+        port: The port of the Ollama instance
         
     Returns:
-        Updated instance with model information
+        A dictionary with IP, port, models, and status information
     """
-    ip = instance["ip"]
-    port = instance["port"]
     url = f"http://{ip}:{port}/api/tags"
-    
-    updated_instance = instance.copy()
+    models = []
+    status = "success"
     
     try:
-        response = requests.get(url, timeout=timeout)
+        response = requests.get(url, timeout=5)
         response.raise_for_status()
         
         data = response.json()
-        
         if "models" in data:
-            models_list = []
-            for model in data["models"]:
-                # Extract model details
+            for model_data in data["models"]:
+                details = model_data.get("details", {})
                 model_info = {
-                    "name": model.get("name", ""),
-                    "family": model.get("details", {}).get("family", ""),
-                    "parameter_size": model.get("details", {}).get("parameter_size", ""),
-                    "quantization_level": model.get("details", {}).get("quantization_level", ""),
-                    "digest": model.get("digest", ""),
-                    "modified_at": model.get("modified_at", ""),
-                    "size": model.get("size", 0)
+                    "name": model_data.get("name", ""),
+                    "family": details.get("family", ""),
+                    "parameter_size": details.get("parameter_size", ""),
+                    "quantization_level": details.get("quantization_level", ""),
+                    "digest": model_data.get("digest", ""),
+                    "modified_at": model_data.get("modified_at", ""),
+                    "size": model_data.get("size", 0)
                 }
-                models_list.append(model_info)
-            
-            updated_instance["models"] = models_list
-            logger.info(f"Successfully extracted {len(models_list)} models from {ip}:{port}")
-        else:
-            logger.warning(f"No models found in response from {ip}:{port}")
-            updated_instance["models"] = []
-            
+                models.append(model_info)
     except requests.exceptions.RequestException as e:
-        logger.error(f"Network error for {ip}:{port}: {e}")
-        updated_instance["models"] = []
+        logger.error(f"Network error when checking {ip}:{port}: {e}")
+        status = "connection failed"
     except ValueError as e:
         logger.error(f"Invalid JSON from {ip}:{port}: {e}")
-        updated_instance["models"] = []
+        status = "invalid json"
     except Exception as e:
-        logger.exception(f"Unexpected error for {ip}:{port}: {e}")
-        updated_instance["models"] = []
+        logger.exception(f"Unexpected error when checking {ip}:{port}")
+        status = "unexpected error"
     
-    return updated_instance
+    return {
+        "ip": ip,
+        "port": port,
+        "models": models,
+        "status": status
+    }
 
-# Authentication functions
-def verify_password(password: str, stored_password: str) -> bool:
+async def check_ollama_endpoints(entries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     """
-    Verify if the entered password matches the stored password.
+    Check multiple Ollama endpoints concurrently.
     
     Args:
-        password: The entered password
-        stored_password: The stored password (hashed or plaintext)
+        entries: List of entries containing IP and port information
         
     Returns:
-        True if passwords match, False otherwise
+        List of updated entries with model information
     """
-    # Sanitize input
-    password = sanitize_input(password)
+    tasks = []
+    for entry in entries:
+        task = asyncio.create_task(check_ollama_endpoint(entry["ip"], entry["port"]))
+        tasks.append((entry, task))
     
-    # Check if stored password is a bcrypt hash
-    if stored_password.startswith('$2b$'):
-        return bcrypt.checkpw(password.encode('utf-8'), stored_password.encode('utf-8'))
-    else:
-        # Direct comparison for development/testing
-        return password == stored_password
+    results = []
+    for entry, task in tasks:
+        try:
+            result = await task
+            # Merge the result with the original entry
+            # This preserves fields like country, region, and org
+            updated_entry = entry.copy()
+            updated_entry["models"] = result["models"]
+            updated_entry["status"] = result["status"]
+            results.append(updated_entry)
+        except Exception as e:
+            logger.error(f"Error checking endpoint {entry.get('ip')}:{entry.get('port')}: {e}")
+            entry["models"] = []
+            entry["status"] = "error"
+            results.append(entry)
+    
+    return results
 
-# UI creation
-def create_ui():
+# Shodan scanning
+def run_shodan_scan() -> List[Dict[str, Any]]:
     """
-    Create the Gradio UI for the application.
+    Run a Shodan scan for Ollama instances.
     
     Returns:
-        Gradio interface
+        List of entries containing IP, port, and location information
     """
-    # Get secrets
-    admin_password = get_secret("ADMIN_PASSWORD")
-    if not admin_password:
-        admin_password = "admin"  # Default for development (should be replaced in production)
-        logger.warning("Admin password not set, using default (insecure)")
-    
-    shodan_api_key = get_secret("SHODAN_API_KEY")
-    if not shodan_api_key:
-        logger.warning("Shodan API key not set, scans will not work")
-    
-    shodan_query = get_secret("SHODAN_QUERY")
-    if not shodan_query:
-        shodan_query = "product:Ollama port:11434"
-        logger.info(f"Using default Shodan query: {shodan_query}")
-    
-    hf_token = get_secret("HF_TOKEN")
-    
-    # Load dataset
-    dataset_repo_id = "latterworks/llama_checker_results"
-    dataset = get_or_create_dataset(dataset_repo_id, token=hf_token)
-    
-    # Function to search and display models
-    def search_models(family, param_size, name, current_dataset):
-        # Sanitize inputs
-        name = sanitize_input(name)
-        
-        results = []
+    api_key = os.environ.get("SHODAN_API_KEY")
+    if not api_key:
+        logger.error("SHODAN_API_KEY environment variable not set")
+        return []
+    
+    query = os.environ.get("SHODAN_QUERY", "product:Ollama port:11434")
+    
+    try:
+        api = shodan.Shodan(api_key)
+        results = api.search(query, limit=1000)
         
-        for item in current_dataset:
-            if item["models"]:
-                for model in item["models"]:
-                    # Apply filters
-                    if family != "All" and model["family"] != family:
-                        continue
-                    
-                    if param_size != "All" and model["parameter_size"] != param_size:
-                        continue
-                    
-                    if name and name.lower() not in model["name"].lower():
-                        continue
-                    
-                    # Calculate size in GB
-                    size_gb = round(model["size"] / (1024 * 1024 * 1024), 2) if model["size"] else 0
-                    
-                    # Add to results
-                    results.append([
-                        model["name"],
-                        model["family"],
-                        model["parameter_size"],
-                        model["quantization_level"],
-                        size_gb
-                    ])
+        entries = []
+        for result in results["matches"]:
+            entry = {
+                "ip": result.get("ip_str", ""),
+                "port": result.get("port", 0),
+                "country": result.get("location", {}).get("country_name", ""),
+                "region": result.get("location", {}).get("region_name", ""),
+                "org": result.get("org", ""),
+                "models": []
+            }
+            entries.append(entry)
         
-        return results
+        return entries
+    except shodan.APIError as e:
+        logger.error(f"Shodan API error: {e}")
+        return []
+    except Exception as e:
+        logger.exception(f"Unexpected error in run_shodan_scan")
+        return []
+
+# Password validation
+def validate_admin_password(password: str) -> bool:
+    """
+    Validate the admin password.
     
-    # Function to display model details
-    def show_model_details(evt: gr.SelectData, results, is_admin, current_dataset):
-        selected_row = evt.index[0]
-        model_name = results[selected_row][0]
+    Args:
+        password: The entered password to validate
         
-        # Find the model
-        for item in current_dataset:
-            if item["models"]:
-                for model in item["models"]:
-                    if model["name"] == model_name:
-                        details = {
-                            "name": model["name"],
-                            "family": model["family"],
-                            "parameter_size": model["parameter_size"],
-                            "quantization_level": model["quantization_level"],
-                            "digest": model["digest"],
-                            "modified_at": model["modified_at"],
-                            "size_bytes": model["size"],
-                            "size_gb": round(model["size"] / (1024 * 1024 * 1024), 2) if model["size"] else 0
-                        }
-                        
-                        # Include IP and port for admin users
-                        if is_admin:
-                            details["ip"] = item["ip"]
-                            details["port"] = item["port"]
-                            details["country"] = item["country"]
-                            details["region"] = item["region"]
-                            details["org"] = item["org"]
-                        
-                        return details
+    Returns:
+        True if the password is valid, False otherwise
+    """
+    stored_password = os.environ.get("ADMIN_PASSWORD")
+    if not stored_password:
+        logger.error("ADMIN_PASSWORD environment variable not set")
+        return False
+    
+    # If the stored password starts with '$2b', it's a bcrypt hash
+    if stored_password.startswith('$2b'):
+        return bcrypt.checkpw(password.encode('utf-8'), stored_password.encode('utf-8'))
+    else:
+        # Otherwise, do a direct comparison
+        return password == stored_password
+
+# Gradio UI functions
+def get_model_families_and_sizes(dataset: Dataset) -> Tuple[List[str], List[str]]:
+    """
+    Extract all unique model families and parameter sizes from the dataset.
+    
+    Args:
+        dataset: The dataset to extract from
         
-        return {"error": "Model not found"}
+    Returns:
+        Tuple of (families, parameter_sizes)
+    """
+    if dataset is None:
+        return [], []
     
-    # Function to update admin visibility
-    def update_admin_visibility(is_admin):
-        return (
-            gr.update(visible=not is_admin),  # admin_required
-            gr.update(visible=is_admin)       # scan_group
-        )
+    families = set()
+    parameter_sizes = set()
+    
+    for i in range(len(dataset)):
+        models = dataset[i]["models"]
+        if models:
+            for model in models:
+                family = model.get("family")
+                param_size = model.get("parameter_size")
+                
+                if family:
+                    families.add(family)
+                if param_size:
+                    parameter_sizes.add(param_size)
     
-    # Function to perform scan
-    def perform_scan(max_results, is_admin, current_dataset):
-        if not is_admin:
-            return "⚠️ Admin login required", [], current_dataset
+    return sorted(list(families)), sorted(list(parameter_sizes))
+
+def search_models(family: str, parameter_size: str, name: str, dataset: Dataset, is_admin: bool) -> Tuple[List[Dict], Dict]:
+    """
+    Search for models in the dataset based on filters.
+    
+    Args:
+        family: Filter by model family
+        parameter_size: Filter by parameter size
+        name: Filter by model name
+        dataset: The dataset to search in
+        is_admin: Whether the user is an admin
         
-        # Start scan
-        yield "🔍 Starting Shodan scan...", [], current_dataset
+    Returns:
+        Tuple of (filtered_models, empty_details)
+    """
+    if dataset is None:
+        return [], {}
+    
+    # Collect all models from the dataset
+    all_models = []
+    for i in range(len(dataset)):
+        ip = dataset[i]["ip"]
+        port = dataset[i]["port"]
+        models = dataset[i]["models"]
         
-        try:
-            # Get instances from Shodan
-            instances = scan_with_shodan(shodan_api_key, shodan_query, max_results)
-            yield f"🔍 Found {len(instances)} instances. Checking endpoints...", [], current_dataset
-            
-            # Check endpoints using executor
-            updated_instances = []
-            with ThreadPoolExecutor(max_workers=10) as executor:
-                # Create future tasks
-                future_to_instance = {
-                    executor.submit(check_ollama_endpoint, instance): instance
-                    for instance in instances
-                }
+        if models:
+            for model in models:
+                model_copy = model.copy()
                 
-                # Process completed tasks
-                for future in as_completed(future_to_instance):
-                    try:
-                        updated_instance = future.result()
-                        updated_instances.append(updated_instance)
-                    except Exception as e:
-                        instance = future_to_instance[future]
-                        logger.exception(f"Error processing {instance['ip']}:{instance['port']}: {e}")
-                        # In case of error, append the original instance without model info
-                        instance["models"] = []
-                        updated_instances.append(instance)
-            
-            # Update dataset
-            updated_dataset = current_dataset
-            for instance in updated_instances:
-                updated_dataset = update_dataset_entry(updated_dataset, instance)
-            
-            # Push to hub
-            push_dataset_to_hub(updated_dataset, dataset_repo_id, token=hf_token)
-            
-            # Prepare results for display
-            results = []
-            total_models = 0
-            for instance in updated_instances:
-                models_count = len(instance["models"]) if instance["models"] else 0
-                total_models += models_count
-                results.append([
-                    instance["ip"],
-                    instance["port"],
-                    instance["country"],
-                    instance["region"],
-                    instance["org"],
-                    models_count
-                ])
-            
-            yield f"✅ Scan completed! Found {len(instances)} instances with a total of {total_models} models.", results, updated_dataset
-        
-        except Exception as e:
-            logger.exception(f"Error during scan: {e}")
-            yield f"❌ Error during scan: {str(e)}", [], current_dataset
+                # Add source info if admin
+                if is_admin:
+                    model_copy["source_ip"] = ip
+                    model_copy["source_port"] = port
+                
+                # Calculate size in GB
+                if "size" in model:
+                    model_copy["size_gb"] = round(model["size"] / (1024**3), 2)
+                else:
+                    model_copy["size_gb"] = 0
+                
+                all_models.append(model_copy)
     
-    # Create Gradio UI
-    with gr.Blocks(title="Ollama Instance Scanner") as app:
-        # State variables
-        current_dataset = gr.State(dataset)
-        is_admin = gr.State(False)
+    # Apply filters
+    filtered_models = all_models
+    if family:
+        filtered_models = [m for m in filtered_models if m.get("family") == family]
+    if parameter_size:
+        filtered_models = [m for m in filtered_models if m.get("parameter_size") == parameter_size]
+    if name:
+        filtered_models = [m for m in filtered_models if name.lower() in m.get("name", "").lower()]
+    
+    return filtered_models, {}
+
+def select_model(evt: gr.SelectData, models: List[Dict]) -> Dict:
+    """
+    Handle model selection from the table.
+    
+    Args:
+        evt: The selection event
+        models: The list of models
         
-        with gr.Tab("Browse Models"):
-            # Filters
-            with gr.Row():
-                with gr.Column():
-                    # Extract unique values for family and parameter_size
-                    families = set()
-                    parameter_sizes = set()
-                    for item in dataset:
-                        if item["models"]:
-                            for model in item["models"]:
-                                if model["family"]:
-                                    families.add(model["family"])
-                                if model["parameter_size"]:
-                                    parameter_sizes.add(model["parameter_size"])
-                    
-                    # Create dropdowns
-                    family_dropdown = gr.Dropdown(
-                        choices=["All"] + sorted(list(families)),
-                        value="All",
-                        label="Model Family"
-                    )
-                    param_size_dropdown = gr.Dropdown(
-                        choices=["All"] + sorted(list(parameter_sizes)),
-                        value="All",
-                        label="Parameter Size"
-                    )
-                    name_search = gr.Textbox(
-                        value="",
-                        label="Model Name Contains"
-                    )
-            
-            # Search button
-            search_btn = gr.Button("Search")
-            
-            # Results dataframe
-            results_df = gr.DataFrame(
-                value=[],
-                headers=["Name", "Family", "Parameter Size", "Quantization", "Size (GB)"],
-                label="Model Results"
-            )
-            
-            # Detailed info JSON viewer
-            model_details = gr.JSON(label="Model Details")
-            
-            # Connect events
-            search_btn.click(
-                search_models,
-                inputs=[family_dropdown, param_size_dropdown, name_search, current_dataset],
-                outputs=[results_df]
+    Returns:
+        The selected model details
+    """
+    if not models or evt.index >= len(models):
+        return {}
+    
+    return models[evt.index]
+
+async def scan_worker() -> str:
+    """
+    Run the complete scan workflow.
+    
+    Returns:
+        Status message
+    """
+    # Run Shodan scan
+    entries = run_shodan_scan()
+    if not entries:
+        return "No Ollama instances found or scan failed"
+    
+    # Check endpoints
+    updated_entries = await check_ollama_endpoints(entries)
+    
+    # Update dataset
+    dataset = get_or_create_dataset()
+    if dataset is not None:
+        update_dataset(dataset, updated_entries)
+        return f"Scan completed. Found {len(entries)} Ollama instances."
+    else:
+        return "Scan completed but failed to update dataset"
+
+# Main application
+def create_app():
+    # Load the dataset
+    dataset = get_or_create_dataset()
+    
+    # Get model families and parameter sizes
+    families, parameter_sizes = [], []
+    if dataset is not None:
+        families, parameter_sizes = get_model_families_and_sizes(dataset)
+    
+    with gr.Blocks(title="Ollama Instance Explorer") as app:
+        # Admin login section
+        with gr.Row():
+            admin_password = gr.Textbox(
+                label="Admin Password",
+                type="password",
+                placeholder="Enter admin password"
             )
-            
-            results_df.select(
-                show_model_details,
-                inputs=[results_df, is_admin, current_dataset],
-                outputs=[model_details]
+            login_button = gr.Button("Login")
+            login_status = gr.Textbox(
+                label="Login Status",
+                value="",
+                interactive=False
             )
         
-        with gr.Tab("Admin Login"):
-            admin_pw_input = gr.Textbox(
+        # Admin state
+        is_admin = gr.State(False)
+        
+        # Admin-only section
+        with gr.Tab("Shodan Scan", visible=False) as admin_tab:
+            scan_button = gr.Button("Start Scan")
+            scan_status = gr.Textbox(
+                label="Scan Status",
                 value="",
-                type="password",
-                label="Admin Password"
+                interactive=False
             )
-            login_btn = gr.Button("Login")
-            login_status = gr.Markdown("Not logged in")
             
-            def admin_login(password):
-                if verify_password(password, admin_password):
-                    return "✅ Successfully logged in as admin", True
-                else:
-                    return "❌ Invalid password", False
+            def on_scan_click():
+                # We can't use async directly with Gradio, so use asyncio.run
+                try:
+                    return asyncio.run(scan_worker())
+                except Exception as e:
+                    logger.exception("Error during scan")
+                    return f"Error during scan: {str(e)}"
             
-            login_btn.click(
-                admin_login,
-                inputs=[admin_pw_input],
-                outputs=[login_status, is_admin]
+            scan_button.click(
+                on_scan_click,
+                inputs=[],
+                outputs=[scan_status]
             )
         
-        with gr.Tab("Shodan Scan") as shodan_tab:
-            # This tab is initially hidden and only shown to admins
-            admin_required = gr.Markdown("⚠️ Admin login required to access this feature")
+        # Public section
+        with gr.Tab("Browse Models"):
+            with gr.Row():
+                family_filter = gr.Dropdown(
+                    label="Family",
+                    choices=[""] + families,
+                    value=""
+                )
+                parameter_size_filter = gr.Dropdown(
+                    label="Parameter Size",
+                    choices=[""] + parameter_sizes,
+                    value=""
+                )
+                name_filter = gr.Textbox(
+                    label="Name Search",
+                    placeholder="Enter model name to search"
+                )
+            
+            search_button = gr.Button("Search")
             
-            with gr.Group(visible=False) as scan_group:
-                max_results = gr.Slider(
-                    minimum=10,
-                    maximum=1000,
-                    value=100,
-                    step=10,
-                    label="Max Results"
+            with gr.Row():
+                models_table = gr.DataFrame(
+                    headers=["name", "family", "parameter_size", "quantization_level", "size_gb"],
+                    datatype=["str", "str", "str", "str", "number"],
+                    interactive=False
                 )
-                scan_btn = gr.Button("Start Scan")
-                scan_status = gr.Markdown("Ready to scan")
+                model_details = gr.JSON(label="Model Details")
+            
+            def on_search(family, parameter_size, name, admin_status):
+                models, _ = search_models(family, parameter_size, name, dataset, admin_status)
                 
-                # Admin results dataframe with IP and port
-                admin_results_df = gr.DataFrame(
-                    value=[],
-                    headers=["IP", "Port", "Country", "Region", "Organization", "Models Count"],
-                    label="Scan Results"
-                )
+                # Create DataFrame-friendly format
+                df_data = []
+                for model in models:
+                    row = {
+                        "name": model.get("name", ""),
+                        "family": model.get("family", ""),
+                        "parameter_size": model.get("parameter_size", ""),
+                        "quantization_level": model.get("quantization_level", ""),
+                        "size_gb": model.get("size_gb", 0)
+                    }
+                    df_data.append(row)
+                
+                return df_data, {}
             
-            # Connect events
-            is_admin.change(
-                update_admin_visibility,
-                inputs=[is_admin],
-                outputs=[admin_required, scan_group]
+            search_button.click(
+                on_search,
+                inputs=[family_filter, parameter_size_filter, name_filter, is_admin],
+                outputs=[models_table, model_details]
             )
             
-            scan_btn.click(
-                perform_scan,
-                inputs=[max_results, is_admin, current_dataset],
-                outputs=[scan_status, admin_results_df, current_dataset]
+            models_table.select(
+                select_model,
+                inputs=[models_table],
+                outputs=[model_details]
             )
+        
+        # Handle login
+        def on_login(password):
+            if validate_admin_password(password):
+                return True, gr.update(visible=True), "Login successful"
+            else:
+                return False, gr.update(visible=False), "Invalid password"
+        
+        login_button.click(
+            on_login,
+            inputs=[admin_password],
+            outputs=[is_admin, admin_tab, login_status]
+        )
+        
+        # Initial search on load
+        app.load(
+            lambda: on_search("", "", "", False),
+            inputs=None,
+            outputs=[models_table, model_details]
+        )
     
     return app
 
-# Main function
-def main():
-    app = create_ui()
-    app.launch()
-
+# Run the app
 if __name__ == "__main__":
-    main()
+    app = create_app()
+    app.launch()