import shodan
import gradio as gr
import logging
from abc import ABC, abstractmethod
from typing import List, Optional
import json
import asyncio

# Logging setup
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

# Custom exceptions
class ScannerException(Exception):
    """Base exception for scanner errors."""
    pass

# Instance class (assuming a basic structure since it wasn’t provided)
class Instance:
    """Represents an Ollama instance."""
    def __init__(self, ip: str, port: int, owner_id: Optional[str] = None):
        self.ip = ip
        self.port = port
        self.owner_id = owner_id

    @classmethod
    def from_shodan(cls, result: dict, owner_id: Optional[str] = None) -> 'Instance':
        """Create Instance from Shodan result."""
        return cls(ip=result['ip_str'], port=result['port'], owner_id=owner_id)

    def to_dict(self) -> dict:
        """Convert to dict for JSON output."""
        return {"ip": self.ip, "port": self.port, "owner_id": self.owner_id}

# Scanner interface
class IScanner(ABC):
    """Scanner interface for finding Ollama instances."""
    @abstractmethod
    async def scan_instances(self, limit: int = 100) -> List[Instance]:
        pass

# ShodanScanner implementation
class ShodanScanner(IScanner):
    """Scanner for finding public Ollama instances using Shodan."""
    def __init__(self, api_key: str, owner_id: Optional[str] = None):
        self.api_key = api_key
        self.owner_id = owner_id

    async def scan_instances(self, limit: int = 100) -> List[Instance]:
        """Scan for Ollama instances using Shodan."""
        try:
            api = shodan.Shodan(self.api_key)
            results = api.search('port:11434 "Ollama"', limit=limit)
            instances = [Instance.from_shodan(result, self.owner_id) for result in results['matches']]
            return instances
        except shodan.APIError as e:
            raise ScannerException(f"Shodan API error: {str(e)}")
        except Exception as e:
            raise ScannerException(f"Error during Shodan scan: {str(e)}")

# Gradio processing function
async def scan_ollama_instances(api_key: str, limit: int, owner_id: str = "") -> str:
    """Process Shodan scan and return results."""
    try:
        if not api_key:
            return "Yo, drop a Shodan API key, fam!"
        
        scanner = ShodanScanner(api_key, owner_id if owner_id else None)
        instances = await scanner.scan_instances(limit=limit)
        
        if not instances:
            return "No Ollama instances found, dawg!"

        output = ["Found Ollama Instances:"]
        for i, instance in enumerate(instances, 1):
            output.append(f"{i}. {json.dumps(instance.to_dict(), indent=2)}")
        
        logger.info(f"Scanned {len(instances)} Ollama instances")
        return "\n\n".join(output)
    
    except ScannerException as e:
        logger.error(f"Scan failed: {str(e)}")
        return f"Scan crashed: {str(e)}"
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}")
        return f"Shit hit the fan, fam: {str(e)}"

# Gradio wrapper to run async function
def gradio_scan(api_key: str, limit: int, owner_id: str) -> str:
    """Wrapper to run async scan in Gradio."""
    return asyncio.run(scan_ollama_instances(api_key, limit, owner_id))

# Gradio interface
demo = gr.Interface(
    fn=gradio_scan,
    inputs=[
        gr.Textbox(label="Shodan API Key", placeholder="Enter your Shodan API key here", type="password"),
        gr.Slider(label="Max Results", minimum=1, maximum=1000, value=100, step=1),
        gr.Textbox(label="Owner ID (Optional)", placeholder="e.g., your_user_id", value="")
    ],
    outputs=gr.Textbox(label="Ollama Instances Found"),
    title="Ollama Instance Scanner",
    description="Scan for public Ollama instances using Shodan, Bay Area style! Drop your API key and let’s roll.",
    allow_flagging="never"
)

if __name__ == "__main__":
    logger.info("Firin’ up the Ollama scanner on Gradio...")
    demo.launch(server_name="0.0.0.0", server_port=7860)