Add webhook secret validation

ai_service.py

@@ -100,8 +100,6 @@ async def get_ai_response(user_query: str, telegram_id: int):
         response_message = completion.choices[0].message
 
     final_response = clean_ai_response(response_message.content if response_message.content else "")
-    print(f"--- AI Raw Response: {repr(response_message.content)} ---")
-    print(f"--- AI Final Response: {repr(final_response)} ---")
 
     if db_manager:
         db_manager.save_message(telegram_id, user_query, "user")

config.py

@@ -10,6 +10,7 @@ load_dotenv()
 PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY")
 HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HF_API_KEY")
 TELEGRAM_TOKEN = os.environ.get("TELEGRAM_TOKEN")
+TELEGRAM_WEBHOOK_SECRET = os.environ.get("TELEGRAM_WEBHOOK_SECRET")
 SUPABASE_URL = os.environ.get("SUPABASE_URL")
 SUPABASE_KEY = os.environ.get("SUPABASE_KEY")
 TELEGRAM_DOMAIN = os.environ.get("TELEGRAM_DOMAIN", "https://api.telegram.org").rstrip("/")

main.py

@@ -1,4 +1,5 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Header, Request
+from security import validate_webhook_secret
 from telegram_handlers import telegram_webhook, WebhookData
 from utils import dns_test, test_ai_response
 
@@ -9,17 +10,27 @@ async def root():
     return {"message": "Hadhramout Bank AI Backend is Live"}
 
 @app.post("/webhook")
-async def webhook(data: WebhookData):
-    return await telegram_webhook(data)
+async def webhook(
+    request: Request,
+    data: WebhookData,
+    x_telegram_bot_api_secret_token: str | None = Header(default=None),
+):
+    validate_webhook_secret(request, x_telegram_bot_api_secret_token)
 
-# @app.post("/test")
-# async def test(data: WebhookData):
-#     return await test_webhook(data)
+    return await telegram_webhook(data)
 
 @app.get("/dns-test")
-async def dns():
+async def dns(
+    request: Request,
+    x_telegram_bot_api_secret_token: str | None = Header(default=None),
+):
+    validate_webhook_secret(request, x_telegram_bot_api_secret_token)
     return await dns_test()
 
 @app.get("/ai-test")
-async def ai():
-    return await test_ai_response()
+async def ai(
+    request: Request,
+    x_telegram_bot_api_secret_token: str | None = Header(default=None),
+):
+    validate_webhook_secret(request, x_telegram_bot_api_secret_token)
+    return await test_ai_response()

security.py

@@ -0,0 +1,54 @@
+from collections import defaultdict, deque
+from time import monotonic
+
+from fastapi import HTTPException, Request
+
+from config import TELEGRAM_WEBHOOK_SECRET
+
+FAILED_ATTEMPT_WINDOW_SECONDS = 60
+FAILED_ATTEMPT_LIMIT = 5
+BLOCK_DURATION_SECONDS = 15 * 60
+
+_failed_secret_attempts: dict[str, deque[float]] = defaultdict(deque)
+_blocked_clients: dict[str, float] = {}
+
+
+def _get_client_key(request: Request) -> str:
+    forwarded_for = request.headers.get("x-forwarded-for")
+    if forwarded_for:
+        return forwarded_for.split(",")[0].strip()
+
+    if request.client and request.client.host:
+        return request.client.host
+
+    return "unknown"
+
+
+def _prune_failed_attempts(client_key: str, now: float) -> deque[float]:
+    attempts = _failed_secret_attempts[client_key]
+    cutoff = now - FAILED_ATTEMPT_WINDOW_SECONDS
+    while attempts and attempts[0] < cutoff:
+        attempts.popleft()
+    return attempts
+
+
+def validate_webhook_secret(request: Request, secret_header: str | None) -> None:
+    if not TELEGRAM_WEBHOOK_SECRET:
+        raise HTTPException(status_code=500, detail="Webhook secret is not configured")
+
+    client_key = _get_client_key(request)
+    now = monotonic()
+    blocked_until = _blocked_clients.get(client_key)
+    if blocked_until and now < blocked_until:
+        raise HTTPException(status_code=429, detail="Too many requests")
+
+    if secret_header != TELEGRAM_WEBHOOK_SECRET:
+        attempts = _prune_failed_attempts(client_key, now)
+        attempts.append(now)
+        if len(attempts) >= FAILED_ATTEMPT_LIMIT:
+            _blocked_clients[client_key] = now + BLOCK_DURATION_SECONDS
+            attempts.clear()
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    _blocked_clients.pop(client_key, None)
+    _failed_secret_attempts.pop(client_key, None)

telegram_handlers.py

@@ -48,8 +48,6 @@ async def telegram_webhook(data: WebhookData):
         username = data.message.chat.username
         first_name = data.message.chat.first_name
 
-        print(f"--- Processing message from {first_name} ---")
-
         if db_manager:
             db_manager.create_or_update_user(telegram_id, username, first_name, data.message.chat.last_name)
 
@@ -75,7 +73,6 @@ async def telegram_webhook(data: WebhookData):
                         "text": final_text if final_text.strip() else ".",
                     }
 
-                    print(f"--- Sending Telegram message (chars={len(payload['text'])}) ---")
 
                     try:                        
                         response = await client.post(TELEGRAM_URL, data=payload)