ads
This commit is contained in:
105
tg-bot/main.py
105
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"""
|
||||
|
||||
@@ -2,5 +2,6 @@ fastapi==0.109.0
|
||||
uvicorn[standard]==0.27.0
|
||||
asyncpg==0.29.0
|
||||
pydantic==2.5.3
|
||||
python-telegram-bot==21.0
|
||||
python-telegram-bot[all]==21.0
|
||||
python-dotenv==1.0.0
|
||||
aiohttp==3.9.1
|
||||
|
||||
Reference in New Issue
Block a user