From 9cc8198ef472fffcee1ec5e5fe9d31205f05cf57 Mon Sep 17 00:00:00 2001 From: ilyastar9999 Date: Thu, 4 Dec 2025 13:32:41 +0300 Subject: [PATCH] asd --- tg-bot/main.py | 133 ++++++++++++++++++++++------------------ tg-bot/requirements.txt | 2 +- 2 files changed, 76 insertions(+), 59 deletions(-) diff --git a/tg-bot/main.py b/tg-bot/main.py index 19fb411..068fb98 100644 --- a/tg-bot/main.py +++ b/tg-bot/main.py @@ -3,13 +3,14 @@ Telegram Bot for A/D Infrastructure Sends notifications to group chat """ import os +import asyncio +import aiohttp from datetime import datetime from fastapi import FastAPI, HTTPException, Depends, Header from pydantic import BaseModel import asyncpg 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,51 +18,81 @@ DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://adctrl:adctrl@postgres:54 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", "") +CONTROLLER_API = os.getenv("CONTROLLER_API", "http://controller:8001") # Database pool and bot db_pool = None bot = None -app_telegram = None +update_offset = 0 +polling_task = None -async def button_handler(update: Update, context: ContextTypes.DEFAULT_TYPE): - """Handle inline button clicks""" - query = update.callback_query - await query.answer() +async def handle_button_click(callback_data: str, chat_id: int, message_id: int): + """Handle inline button click""" + if not callback_data.startswith("service_"): + return - callback_data = query.data - chat_id = query.message.chat_id + parts = callback_data.rsplit("_", 1) + if len(parts) != 2: + return - 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") + action = parts[0].replace("service_", "") + try: + service_id = int(parts[1]) + except ValueError: + return + + try: + async with aiohttp.ClientSession() as session: + api_url = f"{CONTROLLER_API}/services/{service_id}/action" + headers = {"Authorization": f"Bearer {SECRET_TOKEN}"} + data = {"action": action} - 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, timeout=aiohttp.ClientTimeout(total=10)) as resp: + if resp.status == 200: + await bot.edit_message_text( + chat_id=chat_id, + message_id=message_id, + text=f"✅ Service action '{action}' executed successfully" + ) + await log_message(chat_id, f"Action: {action} on service {service_id}", True) + else: + error_text = await resp.text() + await bot.edit_message_text( + chat_id=chat_id, + message_id=message_id, + text=f"❌ Failed: {error_text[:100]}" + ) + await log_message(chat_id, f"Action failed: {action}", False, error_text) + except Exception as e: + await log_message(chat_id, f"Action error: {action}", False, str(e)) + +async def poll_updates(): + """Poll Telegram for updates (button clicks)""" + global update_offset + + if not bot: + return + + while True: + try: + updates = await bot.get_updates(offset=update_offset, timeout=30) + + for update in updates: + update_offset = update.update_id + 1 - 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) + if update.callback_query: + query = update.callback_query + await handle_button_click( + query.data, + query.message.chat_id, + query.message.message_id + ) + await bot.answer_callback_query(query.id) + except asyncio.CancelledError: + break 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)) + print(f"Error in polling: {e}") + await asyncio.sleep(5) class MessageRequest(BaseModel): message: str @@ -104,21 +135,21 @@ 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, app_telegram + global db_pool, bot, polling_task 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() + polling_task = asyncio.create_task(poll_updates()) yield - if app_telegram: - await app_telegram.shutdown() + if polling_task: + polling_task.cancel() + try: + await polling_task + except asyncio.CancelledError: + pass await db_pool.close() @@ -239,20 +270,6 @@ 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""" diff --git a/tg-bot/requirements.txt b/tg-bot/requirements.txt index 9073141..5b64634 100644 --- a/tg-bot/requirements.txt +++ b/tg-bot/requirements.txt @@ -2,6 +2,6 @@ fastapi==0.109.0 uvicorn[standard]==0.27.0 asyncpg==0.29.0 pydantic==2.5.3 -python-telegram-bot[all]==21.0 +python-telegram-bot==21.0 python-dotenv==1.0.0 aiohttp==3.9.1