ads

tg-bot/main.py

@@ -7,8 +7,9 @@ from datetime import datetime
 from fastapi import FastAPI, HTTPException, Depends, Header
 from pydantic import BaseModel
 import asyncpg
-from telegram import Bot
+from telegram import Bot, InlineKeyboardButton, InlineKeyboardMarkup, Update
 from telegram.error import TelegramError
+from telegram.ext import Application, CallbackQueryHandler, ContextTypes
 from contextlib import asynccontextmanager
 
 # Configuration
@@ -17,19 +18,56 @@ SECRET_TOKEN = os.getenv("SECRET_TOKEN", "change-me-in-production")
 TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "")
 TELEGRAM_CHAT_ID = os.getenv("TELEGRAM_CHAT_ID", "")
 
-if not TELEGRAM_BOT_TOKEN:
-    print("WARNING: TELEGRAM_BOT_TOKEN not set!")
-
-if not TELEGRAM_CHAT_ID:
-    print("WARNING: TELEGRAM_CHAT_ID not set!")
-
 # Database pool and bot
 db_pool = None
 bot = None
+app_telegram = None
+
+async def button_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
+    """Handle inline button clicks"""
+    query = update.callback_query
+    await query.answer()
+    
+    callback_data = query.data
+    chat_id = query.message.chat_id
+    
+    if callback_data.startswith("service_"):
+        action, service_id = callback_data.rsplit("_", 1)
+        action = action.replace("service_", "")
+        
+        try:
+            import aiohttp
+            controller_url = os.getenv("CONTROLLER_API", "http://controller:8001")
+            
+            async with aiohttp.ClientSession() as session:
+                api_url = f"{controller_url}/services/{service_id}/action"
+                headers = {"Authorization": f"Bearer {SECRET_TOKEN}"}
+                data = {"action": action}
+                
+                async with session.post(api_url, json=data, headers=headers) as resp:
+                    if resp.status == 200:
+                        result = await resp.json()
+                        await query.edit_message_text(
+                            text=f"✅ Service action '{action}' executed successfully\n{query.message.text}"
+                        )
+                        await log_message(chat_id, f"Button action: {action} service {service_id}", True)
+                    else:
+                        error_text = await resp.text()
+                        await query.edit_message_text(
+                            text=f"❌ Failed to execute action: {error_text}\n{query.message.text}"
+                        )
+                        await log_message(chat_id, f"Button action failed: {action} service {service_id}", False, error_text)
+        except Exception as e:
+            await query.edit_message_text(
+                text=f"❌ Error: {str(e)}\n{query.message.text}"
+            )
+            await log_message(chat_id, f"Button action error: {callback_data}", False, str(e))
 
 class MessageRequest(BaseModel):
     message: str
-    chat_id: str = None  # Optional, uses default if not provided
+    chat_id: str = None
+    service_id: int = None
+    service_name: str = None
 
 class BulkMessageRequest(BaseModel):
     messages: list[str]
@@ -66,14 +104,22 @@ async def log_message(chat_id: int, message: str, success: bool, error_message:
 # Lifespan context
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    global db_pool, bot
+    global db_pool, bot, app_telegram
     db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
     
     if TELEGRAM_BOT_TOKEN:
         bot = Bot(token=TELEGRAM_BOT_TOKEN)
+        
+        app_telegram = Application.builder().token(TELEGRAM_BOT_TOKEN).build()
+        app_telegram.add_handler(CallbackQueryHandler(button_handler))
+        
+        await app_telegram.initialize()
     
     yield
     
+    if app_telegram:
+        await app_telegram.shutdown()
+    
     await db_pool.close()
 
 app = FastAPI(title="Telegram Bot API", lifespan=lifespan)
@@ -89,7 +135,7 @@ async def health_check():
 
 @app.post("/send", dependencies=[Depends(verify_token)])
 async def send_message(request: MessageRequest):
-    """Send a message to telegram chat"""
+    """Send a message to telegram chat with optional service control buttons"""
     if not bot:
         raise HTTPException(status_code=503, detail="Telegram bot not configured")
     
@@ -98,14 +144,24 @@ async def send_message(request: MessageRequest):
         raise HTTPException(status_code=400, detail="No chat_id provided and no default configured")
     
     try:
-        # Send message
-        message = await bot.send_message(
-            chat_id=int(chat_id),
-            text=request.message,
-            parse_mode='HTML'
-        )
+        kwargs = {
+            "chat_id": int(chat_id),
+            "text": request.message,
+            "parse_mode": "HTML"
+        }
         
-        # Log success
+        # Add inline buttons for service control if service_id is provided
+        if request.service_id and request.service_name:
+            keyboard = [
+                [
+                    InlineKeyboardButton("▶️ Start", callback_data=f"service_start_{request.service_id}"),
+                    InlineKeyboardButton("⏹️ Stop", callback_data=f"service_stop_{request.service_id}"),
+                    InlineKeyboardButton("🔄 Restart", callback_data=f"service_restart_{request.service_id}")
+                ]
+            ]
+            kwargs["reply_markup"] = InlineKeyboardMarkup(keyboard)
+        
+        message = await bot.send_message(**kwargs)
         await log_message(int(chat_id), request.message, True)
         
         return {
@@ -115,7 +171,6 @@ async def send_message(request: MessageRequest):
         }
     
     except TelegramError as e:
-        # Log failure
         await log_message(int(chat_id), request.message, False, str(e))
         raise HTTPException(status_code=500, detail=f"Failed to send message: {str(e)}")
 
@@ -184,6 +239,20 @@ async def get_stats():
     finally:
         await release_db(conn)
 
+@app.post("/webhook")
+async def webhook(update_data: dict):
+    """Handle Telegram webhook updates for button callbacks"""
+    if not app_telegram:
+        raise HTTPException(status_code=503, detail="Telegram app not configured")
+    
+    try:
+        update = Update.de_json(update_data, bot)
+        if update:
+            await app_telegram.process_update(update)
+        return {"status": "ok"}
+    except Exception as e:
+        return {"status": "error", "message": str(e)}
+
 @app.post("/test", dependencies=[Depends(verify_token)])
 async def test_connection():
     """Test telegram bot connection"""