asd
This commit is contained in:
129
tg-bot/main.py
129
tg-bot/main.py
@@ -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"""
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user