Update app.py

app.py

@@ -1,532 +1,421 @@
 import os
 import logging
 import asyncio
-import bcrypt
-import requests
-import shodan
+import time
+from typing import Dict, List, Optional, Any, Tuple
+
 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
+import datasets
+import shodan
+import requests
 
-# Setup logging
+# Set up logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
 
-# 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(dataset_name: str = "latterworks/llama_checker_results") -> Optional[Dataset]:
+def load_or_create_dataset():
     """
-    Load the dataset or create it if it doesn't exist.
+    Load or create the dataset.
     
-    Args:
-        dataset_name: The name of the dataset on Hugging Face Hub
-        
     Returns:
-        The dataset or None if there was an error
+        HuggingFace dataset
     """
+    hf_token = os.getenv("HF_TOKEN")
+    if not hf_token:
+        raise ValueError("HF_TOKEN environment variable is not set")
+    
     try:
-        # Try to load the dataset
-        token = os.environ.get("HF_TOKEN")
-        
-        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
+        dataset = datasets.load_dataset("latterworks/llama_checker_results", use_auth_token=hf_token)
+        # Convert to in-memory dataset for easier manipulation
+        dataset = dataset['train']
+    except FileNotFoundError:
+        # Dataset doesn't exist, create it
+        dataset = datasets.Dataset.from_dict({"ip": [], "port": [], "country": [], "region": [], "org": [], "models": []})
+        dataset.push_to_hub("latterworks/llama_checker_results", token=hf_token)
+        dataset = datasets.load_dataset("latterworks/llama_checker_results", use_auth_token=hf_token)['train']
     except Exception as e:
-        logger.exception(f"Unexpected error in get_or_create_dataset: {e}")
-        return None
+        logging.error(f"Failed to load or create dataset: {e}")
+        raise  # Re-raise the exception to stop the application
+    
+    return dataset
+
 
-def update_dataset(dataset: Dataset, new_entries: List[Dict[str, Any]]) -> Optional[Dataset]:
+def scan_shodan(progress=gr.Progress()) -> List[Dict]:
     """
-    Update the dataset with new entries from Shodan scan.
+    Scan Shodan for Ollama instances.
     
     Args:
-        dataset: The dataset to update
-        new_entries: List of new entries to add or update in the dataset
-        
+        progress: Gradio progress bar
+    
     Returns:
-        The updated dataset or None if there was an error
+        List of dictionaries containing information about Ollama instances
     """
-    if dataset is None:
-        logger.error("Cannot update None dataset")
-        return None
+    # Validate Shodan API key exists
+    shodan_api_key = os.getenv("SHODAN_API_KEY")
+    if not shodan_api_key:
+        raise ValueError("SHODAN_API_KEY environment variable is not set")
+    
+    # Get Shodan query
+    shodan_query = os.getenv("SHODAN_QUERY", "product:Ollama port:11434")
     
-    # Convert dataset to dictionaries for easier manipulation
-    dataset_dict = dataset.to_dict()
+    # Initialize Shodan API
+    api = shodan.Shodan(shodan_api_key)
     
-    # Add new entries
-    for entry in new_entries:
-        ip = entry.get("ip")
-        port = entry.get("port")
+    try:
+        # Search Shodan
+        logger.info(f"Searching Shodan with query: {shodan_query}")
+        results = api.search(shodan_query)
+        
+        # Process results
+        instances = []
+        total_results = results['total']
+        logger.info(f"Found {total_results} results")
         
-        # 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
+        # Set up progress bar
+        progress(0, desc="Scanning Shodan for Ollama instances")
         
-        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)
+        for i, result in enumerate(results['matches']):
+            progress((i+1)/len(results['matches']), desc=f"Processing result {i+1}/{len(results['matches'])}")
+            
+            instance = {
+                'ip': result['ip_str'],
+                'port': result.get('port', 11434),
+                'country': result.get('location', {}).get('country_name'),
+                'region': result.get('location', {}).get('region_name'),
+                'org': result.get('org'),
+                'models': []
+            }
+            instances.append(instance)
+            
+        return instances
     
-    return updated_dataset
+    except shodan.APIError as e:
+        logger.error(f"Shodan API error: {e}")
+        raise
+    except Exception as e:
+        logger.error(f"Error during Shodan scan: {e}")
+        raise
+
 
-# Ollama endpoint checking
-async def check_ollama_endpoint(ip: str, port: int) -> Dict[str, Any]:
+async def check_single_endpoint(ip: str, port: int) -> Optional[List[Dict]]:
     """
-    Check a single Ollama endpoint and retrieve model information.
+    Check a single Ollama endpoint for available models.
     
     Args:
-        ip: The IP address of the Ollama instance
-        port: The port of the Ollama instance
-        
+        ip: IP address of the endpoint
+        port: Port number of the endpoint
+    
     Returns:
-        A dictionary with IP, port, models, and status information
+        List of models if successful, None otherwise
     """
     url = f"http://{ip}:{port}/api/tags"
-    models = []
-    status = "success"
     
     try:
+        # Set a timeout of 5 seconds
         response = requests.get(url, timeout=5)
-        response.raise_for_status()
-        
-        data = response.json()
-        if "models" in data:
-            for model_data in data["models"]:
-                details = model_data.get("details", {})
-                model_info = {
-                    "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.append(model_info)
+        if response.status_code == 200:
+            data = response.json()
+            return data.get('models', [])
+        else:
+            logger.warning(f"Failed to get models from {ip}:{port}, status code: {response.status_code}")
+            return None
     except requests.exceptions.RequestException as e:
-        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}")
-        status = "invalid json"
+        logger.warning(f"Error connecting to {ip}:{port}: {e}")
+        return None
     except Exception as e:
-        logger.exception(f"Unexpected error when checking {ip}:{port}")
-        status = "unexpected error"
-    
-    return {
-        "ip": ip,
-        "port": port,
-        "models": models,
-        "status": status
-    }
+        logger.warning(f"Unexpected error checking {ip}:{port}: {e}")
+        return None
 
-async def check_ollama_endpoints(entries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+
+async def check_ollama_endpoints(instances: List[Dict], dataset, progress=gr.Progress()) -> datasets.Dataset:
     """
-    Check multiple Ollama endpoints concurrently.
+    Check all Ollama endpoints for available models asynchronously.
     
     Args:
-        entries: List of entries containing IP and port information
+        instances: List of dictionaries containing information about Ollama instances
+        dataset: HuggingFace dataset
+        progress: Gradio progress bar
         
     Returns:
-        List of updated entries with model information
-    """
-    tasks = []
-    for entry in entries:
-        task = asyncio.create_task(check_ollama_endpoint(entry["ip"], entry["port"]))
-        tasks.append((entry, task))
-    
-    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
-
-# Shodan scanning
-def run_shodan_scan() -> List[Dict[str, Any]]:
-    """
-    Run a Shodan scan for Ollama instances.
-    
-    Returns:
-        List of entries containing IP, port, and location information
+        Updated dataset
     """
-    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)
+    # Validate HF token exists
+    hf_token = os.getenv("HF_TOKEN")
+    if not hf_token:
+        raise ValueError("HF_TOKEN environment variable is not set")
+    
+    # Convert dataset to dictionary for easier manipulation
+    dataset_dict = {f"{item['ip']}:{item['port']}": item for item in dataset}
+    
+    # Process each instance
+    progress(0, desc="Checking Ollama endpoints")
+    for i, instance in enumerate(instances):
+        progress((i+1)/len(instances), desc=f"Checking endpoint {i+1}/{len(instances)}")
         
-        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)
+        ip = instance['ip']
+        port = instance['port']
+        key = f"{ip}:{port}"
         
-        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.
-    
-    Args:
-        password: The entered password to validate
+        # Get models from the endpoint
+        models = await check_single_endpoint(ip, port)
         
-    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
+        if models:
+            # Update instance with models
+            instance['models'] = models
+            
+            # Update or add to dataset dictionary
+            dataset_dict[key] = instance
+    
+    # Convert back to dataset
+    updated_dataset = datasets.Dataset.from_dict({
+        "ip": [item['ip'] for item in dataset_dict.values()],
+        "port": [item['port'] for item in dataset_dict.values()],
+        "country": [item.get('country', '') for item in dataset_dict.values()],
+        "region": [item.get('region', '') for item in dataset_dict.values()],
+        "org": [item.get('org', '') for item in dataset_dict.values()],
+        "models": [item.get('models', []) for item in dataset_dict.values()]
+    })
+    
+    # Push updates to hub
+    updated_dataset.push_to_hub("latterworks/llama_checker_results", token=hf_token)
+    
+    return updated_dataset
+
 
-# Gradio UI functions
-def get_model_families_and_sizes(dataset: Dataset) -> Tuple[List[str], List[str]]:
+def get_unique_values(dataset) -> Tuple[List[str], List[str], List[str]]:
     """
-    Extract all unique model families and parameter sizes from the dataset.
+    Get unique values for family, parameter_size, and name.
     
     Args:
-        dataset: The dataset to extract from
-        
+        dataset: HuggingFace dataset
+    
     Returns:
-        Tuple of (families, parameter_sizes)
+        Tuple of lists containing unique values for family, parameter_size, and name
     """
-    if dataset is None:
-        return [], []
-    
     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)
-    
-    return sorted(list(families)), sorted(list(parameter_sizes))
+    names = set()
+    
+    for item in dataset:
+        for model in item.get('models', []):
+            if 'family' in model and model['family']:
+                families.add(model['family'])
+            if 'parameter_size' in model and model['parameter_size']:
+                parameter_sizes.add(model['parameter_size'])
+            if 'name' in model and model['name']:
+                names.add(model['name'])
+    
+    # Convert to sorted lists and add empty option
+    families = [''] + sorted(list(families))
+    parameter_sizes = [''] + sorted(list(parameter_sizes))
+    names = sorted(list(names))
+    
+    return families, parameter_sizes, names
 
-def search_models(family: str, parameter_size: str, name: str, dataset: Dataset, is_admin: bool) -> Tuple[List[Dict], Dict]:
+
+def search_models(
+    dataset,
+    family: str = "",
+    parameter_size: str = "",
+    name: str = "",
+    is_admin: bool = False
+) -> Tuple[List[Dict], List[Dict]]:
     """
-    Search for models in the dataset based on filters.
+    Search models based on criteria.
     
     Args:
+        dataset: HuggingFace dataset
         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
-        
+    
     Returns:
-        Tuple of (filtered_models, empty_details)
+        Tuple of (results, selected model info)
     """
-    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"]
-        
-        if models:
-            for model in models:
-                model_copy = model.copy()
-                
-                # 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)
-    
-    # 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, {}
+    results = []
+    
+    for item in dataset:
+        for model in item.get('models', []):
+            # Apply filters
+            if family and model.get('family', '') != family:
+                continue
+            if parameter_size and model.get('parameter_size', '') != parameter_size:
+                continue
+            if name and name.lower() not in model.get('name', '').lower():
+                continue
+            
+            # Create result with model info
+            result = {
+                'name': model.get('name', ''),
+                'family': model.get('family', ''),
+                'parameter_size': model.get('parameter_size', ''),
+                'quantization_level': model.get('quantization_level', ''),
+                'size': round(model.get('size', 0) / (1024**3), 2)  # Convert to GB
+            }
+            
+            # Add IP and port only for admin users
+            if is_admin:
+                result['ip'] = item['ip']
+                result['port'] = item['port']
+            
+            results.append(result)
+    
+    # For empty result, return empty JSON info
+    selected_model_info = [{}]
+    
+    return results, selected_model_info
+
 
-def select_model(evt: gr.SelectData, models: List[Dict]) -> Dict:
+def get_model_info(model_row: Dict) -> Dict:
     """
-    Handle model selection from the table.
+    Get detailed information about a selected model.
     
     Args:
-        evt: The selection event
-        models: The list of models
-        
+        model_row: Selected model row from the results
+    
     Returns:
-        The selected model details
+        Dictionary containing detailed model information
     """
-    if not models or evt.index >= len(models):
-        return {}
-    
-    return models[evt.index]
+    return model_row
 
-async def scan_worker() -> str:
+
+def create_interface():
     """
-    Run the complete scan workflow.
+    Create Gradio interface for the application.
     
     Returns:
-        Status message
+        Gradio interface
     """
-    # 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"
-            )
-            login_button = gr.Button("Login")
-            login_status = gr.Textbox(
-                label="Login Status",
-                value="",
-                interactive=False
-            )
+    # Load or create dataset
+    dataset = load_or_create_dataset()
+    
+    # Check for admin mode
+    is_admin = os.getenv("ADMIN_MODE", "false").lower() == "true"
+    
+    # Get unique values for dropdown menus
+    families, parameter_sizes, names = get_unique_values(dataset)
+    
+    # Get initial search results
+    initial_results, initial_model_info = search_models(dataset, is_admin=is_admin)
+    
+    # Function to run Shodan scan
+    def run_shodan_scan(progress=gr.Progress()):
+        nonlocal dataset
+        instances = scan_shodan(progress)
+        dataset = asyncio.run(check_ollama_endpoints(instances, dataset, progress))
         
-        # Admin state
-        is_admin = gr.State(False)
+        # Update unique values
+        updated_families, updated_parameter_sizes, updated_names = get_unique_values(dataset)
         
-        # 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="",
-                interactive=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)}"
-            
-            scan_button.click(
-                on_scan_click,
-                inputs=[],
-                outputs=[scan_status]
-            )
+        # Update search results
+        updated_results, updated_model_info = search_models(dataset, is_admin=is_admin)
+        
+        return (
+            updated_families, updated_parameter_sizes,
+            updated_results, updated_model_info
+        )
+    
+    # Function to run model search
+    def run_search(family, parameter_size, name):
+        results, model_info = search_models(dataset, family, parameter_size, name, is_admin=is_admin)
+        return results, model_info
+    
+    # Function to get model details when a row is selected
+    def select_model(evt: gr.SelectData, results):
+        if evt.index[0] < len(results):
+            selected = results[evt.index[0]]
+            return selected
+        return {}
+    
+    # Create Gradio interface
+    with gr.Blocks(title="Ollama Instance Scanner") as interface:
+        gr.Markdown("# Ollama Instance Scanner")
         
-        # Public section
-        with gr.Tab("Browse Models"):
-            with gr.Row():
-                family_filter = gr.Dropdown(
-                    label="Family",
-                    choices=[""] + families,
-                    value=""
+        with gr.Tabs():
+            # Browse Models tab
+            with gr.TabItem("Browse Models"):
+                with gr.Row():
+                    with gr.Column():
+                        family_dropdown = gr.Dropdown(
+                            choices=families,
+                            label="Model Family",
+                            value=""
+                        )
+                        parameter_size_dropdown = gr.Dropdown(
+                            choices=parameter_sizes,
+                            label="Parameter Size",
+                            value=""
+                        )
+                        name_search = gr.Textbox(
+                            label="Model Name",
+                            placeholder="Search by name..."
+                        )
+                        search_button = gr.Button("Search")
+                
+                results_df = gr.DataFrame(
+                    value=initial_results,
+                    label="Search Results",
+                    headers=["name", "family", "parameter_size", "quantization_level", "size"],
+                    row_count=10,
+                    interactive=False
                 )
-                parameter_size_filter = gr.Dropdown(
-                    label="Parameter Size",
-                    choices=[""] + parameter_sizes,
-                    value=""
+                
+                model_info = gr.JSON(
+                    value=initial_model_info[0] if initial_model_info else {},
+                    label="Model Details"
                 )
-                name_filter = gr.Textbox(
-                    label="Name Search",
-                    placeholder="Enter model name to search"
+                
+                # Event handlers
+                search_button.click(
+                    fn=run_search,
+                    inputs=[family_dropdown, parameter_size_dropdown, name_search],
+                    outputs=[results_df, model_info]
                 )
-            
-            search_button = gr.Button("Search")
-            
-            with gr.Row():
-                models_table = gr.DataFrame(
-                    headers=["name", "family", "parameter_size", "quantization_level", "size_gb"],
-                    datatype=["str", "str", "str", "str", "number"],
-                    interactive=False
+                
+                results_df.select(
+                    fn=select_model,
+                    inputs=[results_df],
+                    outputs=[model_info]
                 )
-                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)
-                
-                # 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)
+            # Shodan Scan tab
+            with gr.TabItem("Shodan Scan"):
+                # Check if Shodan API key is available
+                shodan_api_key = os.getenv("SHODAN_API_KEY")
                 
-                return df_data, {}
-            
-            search_button.click(
-                on_search,
-                inputs=[family_filter, parameter_size_filter, name_filter, is_admin],
-                outputs=[models_table, model_details]
-            )
-            
-            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]
-        )
+                if shodan_api_key:
+                    scan_button = gr.Button("Start Scan")
+                    scan_output = gr.Markdown("Press the button to start scanning Shodan for Ollama instances.")
+                    
+                    # Event handlers
+                    scan_button.click(
+                        fn=run_shodan_scan,
+                        outputs=[
+                            family_dropdown, parameter_size_dropdown,
+                            results_df, model_info
+                        ]
+                    )
+                else:
+                    gr.Markdown("## Shodan API key not configured")
+                    gr.Markdown(
+                        "To use the Shodan scan feature, you need to set the `SHODAN_API_KEY` "
+                        "environment variable in your Hugging Face Space settings."
+                    )
     
-    return app
+    return interface
+
+
+def main():
+    """Main function to run the application."""
+    # Create and launch interface
+    interface = create_interface()
+    interface.launch()
+
 
-# Run the app
 if __name__ == "__main__":
-    app = create_app()
-    app.launch()
+    main()