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
"""
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_", "")
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 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:
import aiohttp
controller_url = os.getenv("CONTROLLER_API", "http://controller:8001")
updates = await bot.get_updates(offset=update_offset, timeout=30)
async with aiohttp.ClientSession() as session:
api_url = f"{controller_url}/services/{service_id}/action"
headers = {"Authorization": f"Bearer {SECRET_TOKEN}"}
data = {"action": action}
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"""

View File

@@ -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