baton/backend/middleware.py
2026-03-21 07:36:33 +02:00

72 lines
2.1 KiB
Python

from __future__ import annotations
import secrets
import time
from typing import Optional
from fastapi import Depends, Header, HTTPException, Request
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from backend import config
_bearer = HTTPBearer(auto_error=False)
_RATE_LIMIT = 5
_RATE_WINDOW = 600 # 10 minutes
_SIGNAL_RATE_LIMIT = 10
_SIGNAL_RATE_WINDOW = 60 # 1 minute
def _get_client_ip(request: Request) -> str:
return (
request.headers.get("X-Real-IP")
or request.headers.get("X-Forwarded-For", "").split(",")[0].strip()
or (request.client.host if request.client else "unknown")
)
async def verify_webhook_secret(
x_telegram_bot_api_secret_token: str = Header(default=""),
) -> None:
if not secrets.compare_digest(
x_telegram_bot_api_secret_token, config.WEBHOOK_SECRET
):
raise HTTPException(status_code=403, detail="Forbidden")
async def verify_admin_token(
credentials: Optional[HTTPAuthorizationCredentials] = Depends(_bearer),
) -> None:
if credentials is None or not secrets.compare_digest(
credentials.credentials, config.ADMIN_TOKEN
):
raise HTTPException(status_code=401, detail="Unauthorized")
async def rate_limit_register(request: Request) -> None:
counters = request.app.state.rate_counters
client_ip = _get_client_ip(request)
now = time.time()
count, window_start = counters.get(client_ip, (0, now))
if now - window_start >= _RATE_WINDOW:
count = 0
window_start = now
count += 1
counters[client_ip] = (count, window_start)
if count > _RATE_LIMIT:
raise HTTPException(status_code=429, detail="Too Many Requests")
async def rate_limit_signal(request: Request) -> None:
counters = request.app.state.rate_counters
key = f"sig:{_get_client_ip(request)}"
now = time.time()
count, window_start = counters.get(key, (0, now))
if now - window_start >= _SIGNAL_RATE_WINDOW:
count = 0
window_start = now
count += 1
counters[key] = (count, window_start)
if count > _SIGNAL_RATE_LIMIT:
raise HTTPException(status_code=429, detail="Too Many Requests")