diff --git a/main.py b/main.py index 5973f89..3a7e992 100644 --- a/main.py +++ b/main.py @@ -9,6 +9,8 @@ import re import httpx import struct import termios +from telegram import Update +from telegram.ext import ApplicationBuilder, ContextTypes, MessageHandler, filters from google import genai from google.genai import types import json @@ -44,6 +46,11 @@ OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o") OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3") NVIDIA_MODEL = os.getenv("NVIDIA_MODEL", "moonshotai/kimi-k2.5") +# Telegram Bot Konfiguration +TELEGRAM_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "") +ALLOWED_ID = os.getenv("ALLOWED_TELEGRAM_USER_ID", "") +telegram_app = None + # --- DATENBANK INITIALISIERUNG (ERWEITERT) --- def init_db(): conn = sqlite3.connect(DB_PATH) @@ -163,6 +170,89 @@ async def get_ai_response(user_input, system_prompt): return ai_msg +async def handle_telegram_message(update: Update, context: ContextTypes.DEFAULT_TYPE): + # Sicherheits-Check: Nur deine Chat-ID wird akzeptiert + if str(update.effective_chat.id) != ALLOWED_ID: + await update.message.reply_text("Zugriff verweigert. 🔒") + return + + user_msg = update.message.text + + # "Tippt..." Status in Telegram anzeigen (schöne UX) + await update.message.reply_chat_action(action="typing") + + # 1. KI Antwort holen + ai_response = await get_ai_response(user_msg, get_system_prompt()) + + # 2. Befehle extrahieren + commands = re.findall(r'(.*?)', ai_response, re.I | re.S) + clean_msg = re.sub(r'.*?', '', ai_response, flags=re.I | re.S).strip() + + # KI Text-Antwort senden + if clean_msg: + await update.message.reply_text(clean_msg) + + # 3. Befehle ausführen und Ergebnisse als Telegram-Nachricht senden + if commands: + for target, cmd in commands: + await update.message.reply_text(f"⏳ Führe aus auf *{target}*:\n`{cmd}`", parse_mode='Markdown') + + # Node in DB suchen + conn = get_db() + n = conn.execute('SELECT * FROM nodes WHERE ip=? OR name=?', (target.strip(), target.strip())).fetchone() + conn.close() + + if n: + # Befehl ausführen und Output abfangen + try: + proc = await asyncio.create_subprocess_shell( + f"ssh -o StrictHostKeyChecking=no {n['user']}@{n['ip']} '{cmd}'", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.STDOUT + ) + stdout, _ = await proc.communicate() + output = stdout.decode('utf-8', errors='ignore').strip() + + # Ergebnis an Telegram senden (abgeschnitten auf 4000 Zeichen, da Telegram Limits hat) + result_text = output[:4000] if output else "✅ Befehl ohne Output ausgeführt." + await update.message.reply_text(f"💻 **Output von {n['name']}:**\n```\n{result_text}\n```", parse_mode='Markdown') + + # Auch ins KI-Gedächtnis schreiben + chat_history.append({"role": "user", "content": f"[SYSTEM] Befehl '{cmd}' auf {target} fertig:\n{result_text}"}) + except Exception as e: + await update.message.reply_text(f"❌ Fehler bei der Ausführung: {e}") + else: + await update.message.reply_text(f"⚠️ Node '{target}' nicht in der Datenbank gefunden.") + +# --- FASTAPI LIFESPAN EVENTS (Bot starten/stoppen) --- +@app.on_event("startup") +async def startup_event(): + global telegram_app + if TELEGRAM_TOKEN and ALLOWED_ID: + print("🤖 Starte Telegram Bot im Hintergrund...") + telegram_app = ApplicationBuilder().token(TELEGRAM_TOKEN).build() + + # Leitet alle Text-Nachrichten an unsere Funktion weiter + telegram_app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_telegram_message)) + + # Bot asynchron in die FastAPI Event-Loop einhängen + await telegram_app.initialize() + await telegram_app.start() + await telegram_app.updater.start_polling() + print("✅ Telegram Bot lauscht!") + else: + print("ℹ️ Telegram Bot inaktiv (Token oder ID fehlen in der .env).") + +@app.on_event("shutdown") +async def shutdown_event(): + global telegram_app + if telegram_app: + print("🛑 Stoppe Telegram Bot...") + await telegram_app.updater.stop() + await telegram_app.stop() + await telegram_app.shutdown() + + # --- WebSocket Manager --- class ConnectionManager: def __init__(self): self.active_connections = []