Files
PiDoBot/main.py
2026-03-03 22:34:25 +00:00

126 lines
4.6 KiB
Python

import os
import subprocess
import sqlite3
import asyncio
from fastapi import FastAPI, WebSocket, BackgroundTasks, Request, Form, WebSocketDisconnect
from fastapi.responses import RedirectResponse
from fastapi.templating import Jinja2Templates
app = FastAPI()
templates = Jinja2Templates(directory="templates")
SSH_KEY = os.path.expanduser("~/.ssh/id_rsa")
DB_PATH = "cluster.db"
# --- WebSocket Manager für Live-Logs ---
class ConnectionManager:
def __init__(self):
self.active_connections: list[WebSocket] = []
async def connect(self, websocket: WebSocket):
await websocket.accept()
self.active_connections.append(websocket)
def disconnect(self, websocket: WebSocket):
self.active_connections.remove(websocket)
async def broadcast(self, message: str):
for connection in self.active_connections:
await connection.send_text(message)
manager = ConnectionManager()
# --- Datenbank & SSH Initialisierung ---
def init_db():
conn = sqlite3.connect(DB_PATH)
conn.execute('CREATE TABLE IF NOT EXISTS nodes (id INTEGER PRIMARY KEY, name TEXT, ip TEXT UNIQUE, user TEXT, status TEXT)')
conn.commit()
conn.close()
def ensure_ssh_key():
if not os.path.exists(SSH_KEY):
subprocess.run(["ssh-keygen", "-t", "rsa", "-N", "", "-f", SSH_KEY], check=True)
init_db()
ensure_ssh_key()
# --- Hintergrund-Task: Key Deployment & Bootstrap ---
async def deploy_and_bootstrap(ip, user, password):
await manager.broadcast(f"Iniziere Setup für {ip}...")
# 1. SSH-Key kopieren
# -o StrictHostKeyChecking=no verhindert die "Are you sure you want to continue connecting" Abfrage
ssh_copy_cmd = f"sshpass -p '{password}' ssh-copy-id -o StrictHostKeyChecking=no -i {SSH_KEY}.pub {user}@{ip}"
process = subprocess.Popen(ssh_copy_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
for line in process.stdout:
await manager.broadcast(f"🔑 SSH: {line.strip()}")
# Datenbank-Status aktualisieren
update_node_status(ip, "Key hinterlegt")
await manager.broadcast(f"✅ SSH-Key erfolgreich auf {ip} kopiert.")
# 2. Abhängigkeiten installieren (Docker & Python)
await manager.broadcast(f"📦 Installiere Docker und Abhängigkeiten auf {ip}...")
bootstrap_cmd = f"ssh {user}@{ip} 'sudo apt-get update && sudo apt-get install -y python3-pip && curl -sSL https://get.docker.com | sh'"
process = subprocess.Popen(bootstrap_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
for line in process.stdout:
# Wir filtern hier ein wenig, um die Konsole nicht zu fluten
if "Progress" not in line:
await manager.broadcast(f"🛠️ {line.strip()[:80]}")
update_node_status(ip, "Online")
await manager.broadcast(f"🏁 Node {ip} ist nun vollständig einsatzbereit!")
def update_node_status(ip, status):
conn = sqlite3.connect(DB_PATH)
conn.execute('UPDATE nodes SET status = ? WHERE ip = ?', (status, ip))
conn.commit()
conn.close()
# --- Routen ---
@app.get("/")
async def index(request: Request):
conn = sqlite3.connect(DB_PATH)
conn.row_factory = sqlite3.Row
nodes = conn.execute('SELECT * FROM nodes').fetchall()
conn.close()
return templates.TemplateResponse("index.html", {"request": request, "nodes": nodes})
@app.post("/add_node")
async def add_node(background_tasks: BackgroundTasks, name: str = Form(...), ip: str = Form(...), user: str = Form(...), password: str = Form(...)):
conn = sqlite3.connect(DB_PATH)
try:
conn.execute('INSERT INTO nodes (name, ip, user, status) VALUES (?, ?, ?, ?)', (name, ip, user, "Setup läuft..."))
conn.commit()
# Starte den SSH-Prozess im Hintergrund
background_tasks.add_task(deploy_and_bootstrap, ip, user, password)
except sqlite3.IntegrityError:
pass
finally:
conn.close()
return RedirectResponse(url="/", status_code=303)
# --- WebSockets ---
@app.websocket("/ws/install_logs")
async def websocket_endpoint(websocket: WebSocket):
await manager.connect(websocket)
try:
while True:
await websocket.receive_text() # Hält die Verbindung offen
except WebSocketDisconnect:
manager.disconnect(websocket)
@app.websocket("/ws/chat")
async def chat_endpoint(websocket: WebSocket):
await websocket.accept()
while True:
msg = await websocket.receive_text()
await websocket.send_text(f"KI: Ich habe '{msg}' empfangen. Aktuell bereite ich die Nodes vor.")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)