Update app.py

app.py

@@ -1,64 +1,113 @@
+import shodan
 import gradio as gr
-from huggingface_hub import InferenceClient
+import logging
+from abc import ABC, abstractmethod
+from typing import List, Optional
+import json
+import asyncio
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+# 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
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+    @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)
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    def to_dict(self) -> dict:
+        """Convert to dict for JSON output."""
+        return {"ip": self.ip, "port": self.port, "owner_id": self.owner_id}
 
-    messages.append({"role": "user", "content": message})
+# Scanner interface
+class IScanner(ABC):
+    """Scanner interface for finding Ollama instances."""
+    @abstractmethod
+    async def scan_instances(self, limit: int = 100) -> List[Instance]:
+        pass
 
-    response = ""
+# 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
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+    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)}")
 
-        response += token
-        yield response
+# 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)}"
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+# 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__":
-    demo.launch()
+    logger.info("Firin’ up the Ollama scanner on Gradio...")
+    demo.launch(server_name="0.0.0.0", server_port=7860)