This commit is contained in:
ilyastar9999
2025-12-04 13:32:41 +03:00
parent 7f841d155d
commit 9cc8198ef4
2 changed files with 76 additions and 59 deletions

View File

@@ -3,13 +3,14 @@ Telegram Bot for A/D Infrastructure
Sends notifications to group chat Sends notifications to group chat
""" """
import os import os
import asyncio
import aiohttp
from datetime import datetime from datetime import datetime
from fastapi import FastAPI, HTTPException, Depends, Header from fastapi import FastAPI, HTTPException, Depends, Header
from pydantic import BaseModel from pydantic import BaseModel
import asyncpg import asyncpg
from telegram import Bot, InlineKeyboardButton, InlineKeyboardMarkup, Update from telegram import Bot, InlineKeyboardButton, InlineKeyboardMarkup, Update
from telegram.error import TelegramError from telegram.error import TelegramError
from telegram.ext import Application, CallbackQueryHandler, ContextTypes
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
# Configuration # 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") SECRET_TOKEN = os.getenv("SECRET_TOKEN", "change-me-in-production")
TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "") TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "")
TELEGRAM_CHAT_ID = os.getenv("TELEGRAM_CHAT_ID", "") TELEGRAM_CHAT_ID = os.getenv("TELEGRAM_CHAT_ID", "")
CONTROLLER_API = os.getenv("CONTROLLER_API", "http://controller:8001")
# Database pool and bot # Database pool and bot
db_pool = None db_pool = None
bot = None bot = None
app_telegram = None update_offset = 0
polling_task = None
async def button_handler(update: Update, context: ContextTypes.DEFAULT_TYPE): async def handle_button_click(callback_data: str, chat_id: int, message_id: int):
"""Handle inline button clicks""" """Handle inline button click"""
query = update.callback_query if not callback_data.startswith("service_"):
await query.answer() return
callback_data = query.data parts = callback_data.rsplit("_", 1)
chat_id = query.message.chat_id if len(parts) != 2:
return
if callback_data.startswith("service_"): action = parts[0].replace("service_", "")
action, service_id = callback_data.rsplit("_", 1) try:
action = action.replace("service_", "") 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 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: try:
import aiohttp updates = await bot.get_updates(offset=update_offset, timeout=30)
controller_url = os.getenv("CONTROLLER_API", "http://controller:8001")
async with aiohttp.ClientSession() as session: for update in updates:
api_url = f"{controller_url}/services/{service_id}/action" update_offset = update.update_id + 1
headers = {"Authorization": f"Bearer {SECRET_TOKEN}"}
data = {"action": action}
async with session.post(api_url, json=data, headers=headers) as resp: if update.callback_query:
if resp.status == 200: query = update.callback_query
result = await resp.json() await handle_button_click(
await query.edit_message_text( query.data,
text=f"✅ Service action '{action}' executed successfully\n{query.message.text}" query.message.chat_id,
) query.message.message_id
await log_message(chat_id, f"Button action: {action} service {service_id}", True) )
else: await bot.answer_callback_query(query.id)
error_text = await resp.text() except asyncio.CancelledError:
await query.edit_message_text( break
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: except Exception as e:
await query.edit_message_text( print(f"Error in polling: {e}")
text=f"❌ Error: {str(e)}\n{query.message.text}" await asyncio.sleep(5)
)
await log_message(chat_id, f"Button action error: {callback_data}", False, str(e))
class MessageRequest(BaseModel): class MessageRequest(BaseModel):
message: str message: str
@@ -104,21 +135,21 @@ async def log_message(chat_id: int, message: str, success: bool, error_message:
# Lifespan context # Lifespan context
@asynccontextmanager @asynccontextmanager
async def lifespan(app: FastAPI): 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) db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
if TELEGRAM_BOT_TOKEN: if TELEGRAM_BOT_TOKEN:
bot = Bot(token=TELEGRAM_BOT_TOKEN) bot = Bot(token=TELEGRAM_BOT_TOKEN)
polling_task = asyncio.create_task(poll_updates())
app_telegram = Application.builder().token(TELEGRAM_BOT_TOKEN).build()
app_telegram.add_handler(CallbackQueryHandler(button_handler))
await app_telegram.initialize()
yield yield
if app_telegram: if polling_task:
await app_telegram.shutdown() polling_task.cancel()
try:
await polling_task
except asyncio.CancelledError:
pass
await db_pool.close() await db_pool.close()
@@ -239,20 +270,6 @@ async def get_stats():
finally: finally:
await release_db(conn) 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)]) @app.post("/test", dependencies=[Depends(verify_token)])
async def test_connection(): async def test_connection():
"""Test telegram bot connection""" """Test telegram bot connection"""

View File

@@ -2,6 +2,6 @@ fastapi==0.109.0
uvicorn[standard]==0.27.0 uvicorn[standard]==0.27.0
asyncpg==0.29.0 asyncpg==0.29.0
pydantic==2.5.3 pydantic==2.5.3
python-telegram-bot[all]==21.0 python-telegram-bot==21.0
python-dotenv==1.0.0 python-dotenv==1.0.0
aiohttp==3.9.1 aiohttp==3.9.1