init

tg-bot/main.py

@@ -0,0 +1,206 @@
+"""
+Telegram Bot for A/D Infrastructure
+Sends notifications to group chat
+"""
+import os
+from datetime import datetime
+from fastapi import FastAPI, HTTPException, Depends, Header
+from pydantic import BaseModel
+import asyncpg
+from telegram import Bot
+from telegram.error import TelegramError
+from contextlib import asynccontextmanager
+
+# Configuration
+DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://adctrl:adctrl@postgres:5432/adctrl")
+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
+
+class MessageRequest(BaseModel):
+    message: str
+    chat_id: str = None  # Optional, uses default if not provided
+
+class BulkMessageRequest(BaseModel):
+    messages: list[str]
+    chat_id: str = None
+
+# Auth dependency
+async def verify_token(authorization: str = Header(None)):
+    if not authorization or not authorization.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Missing or invalid authorization header")
+    
+    token = authorization.replace("Bearer ", "")
+    if token != SECRET_TOKEN:
+        raise HTTPException(status_code=403, detail="Invalid token")
+    return token
+
+# Database functions
+async def get_db():
+    return await db_pool.acquire()
+
+async def release_db(conn):
+    await db_pool.release(conn)
+
+async def log_message(chat_id: int, message: str, success: bool, error_message: str = None):
+    """Log sent message to database"""
+    conn = await db_pool.acquire()
+    try:
+        await conn.execute(
+            "INSERT INTO telegram_messages (chat_id, message, success, error_message) VALUES ($1, $2, $3, $4)",
+            chat_id, message, success, error_message
+        )
+    finally:
+        await db_pool.release(conn)
+
+# Lifespan context
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global db_pool, bot
+    db_pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
+    
+    if TELEGRAM_BOT_TOKEN:
+        bot = Bot(token=TELEGRAM_BOT_TOKEN)
+    
+    yield
+    
+    await db_pool.close()
+
+app = FastAPI(title="Telegram Bot API", lifespan=lifespan)
+
+# API Endpoints
+@app.get("/health")
+async def health_check():
+    return {
+        "status": "ok",
+        "bot_configured": bot is not None,
+        "timestamp": datetime.utcnow().isoformat()
+    }
+
+@app.post("/send", dependencies=[Depends(verify_token)])
+async def send_message(request: MessageRequest):
+    """Send a message to telegram chat"""
+    if not bot:
+        raise HTTPException(status_code=503, detail="Telegram bot not configured")
+    
+    chat_id = request.chat_id or TELEGRAM_CHAT_ID
+    if not chat_id:
+        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'
+        )
+        
+        # Log success
+        await log_message(int(chat_id), request.message, True)
+        
+        return {
+            "status": "sent",
+            "message_id": message.message_id,
+            "chat_id": chat_id
+        }
+    
+    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)}")
+
+@app.post("/send-bulk", dependencies=[Depends(verify_token)])
+async def send_bulk_messages(request: BulkMessageRequest):
+    """Send multiple messages to telegram chat"""
+    if not bot:
+        raise HTTPException(status_code=503, detail="Telegram bot not configured")
+    
+    chat_id = request.chat_id or TELEGRAM_CHAT_ID
+    if not chat_id:
+        raise HTTPException(status_code=400, detail="No chat_id provided and no default configured")
+    
+    results = []
+    for msg in request.messages:
+        try:
+            message = await bot.send_message(
+                chat_id=int(chat_id),
+                text=msg,
+                parse_mode='HTML'
+            )
+            await log_message(int(chat_id), msg, True)
+            results.append({
+                "status": "sent",
+                "message_id": message.message_id,
+                "message": msg[:50] + "..." if len(msg) > 50 else msg
+            })
+        except TelegramError as e:
+            await log_message(int(chat_id), msg, False, str(e))
+            results.append({
+                "status": "failed",
+                "error": str(e),
+                "message": msg[:50] + "..." if len(msg) > 50 else msg
+            })
+    
+    return {"results": results, "total": len(results)}
+
+@app.get("/messages", dependencies=[Depends(verify_token)])
+async def get_message_history(limit: int = 50):
+    """Get message sending history"""
+    conn = await get_db()
+    try:
+        rows = await conn.fetch(
+            "SELECT * FROM telegram_messages ORDER BY sent_at DESC LIMIT $1",
+            limit
+        )
+        return [dict(row) for row in rows]
+    finally:
+        await release_db(conn)
+
+@app.get("/stats", dependencies=[Depends(verify_token)])
+async def get_stats():
+    """Get message statistics"""
+    conn = await get_db()
+    try:
+        total = await conn.fetchval("SELECT COUNT(*) FROM telegram_messages")
+        successful = await conn.fetchval("SELECT COUNT(*) FROM telegram_messages WHERE success = true")
+        failed = await conn.fetchval("SELECT COUNT(*) FROM telegram_messages WHERE success = false")
+        
+        return {
+            "total_messages": total,
+            "successful": successful,
+            "failed": failed,
+            "success_rate": (successful / total * 100) if total > 0 else 0
+        }
+    finally:
+        await release_db(conn)
+
+@app.post("/test", dependencies=[Depends(verify_token)])
+async def test_connection():
+    """Test telegram bot connection"""
+    if not bot:
+        raise HTTPException(status_code=503, detail="Telegram bot not configured")
+    
+    try:
+        me = await bot.get_me()
+        return {
+            "status": "ok",
+            "bot_username": me.username,
+            "bot_name": me.first_name,
+            "bot_id": me.id
+        }
+    except TelegramError as e:
+        raise HTTPException(status_code=500, detail=f"Bot test failed: {str(e)}")
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8003)