Update main.py

main.py

@@ -16,9 +16,9 @@ from slowapi import Limiter, _rate_limit_exceeded_handler
 from slowapi.util import get_remote_address
 from slowapi.errors import RateLimitExceeded
 from slowapi.middleware import SlowAPIMiddleware
-from fastapi import FastAPI, HTTPException, Request, Depends, Header, Security
+from fastapi import FastAPI, HTTPException, Request, Depends, Security
 from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
-from fastapi.security.api_key import APIKeyHeader
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 
 from aiohttp import ClientSession, ClientTimeout, ClientError
@@ -80,24 +80,42 @@ if not API_KEYS or API_KEYS == ['']:
     logger.error("No API keys found. Please set the API_KEYS environment variable.")
     raise Exception("API_KEYS environment variable not set.")
 
-# Define API key security
-API_KEY_NAME = "X-API-Key"
-api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
+# Define API key security using HTTPBearer (Bearer token)
+security = HTTPBearer()
 
-async def get_api_key(api_key: str = Security(api_key_header)):
+async def get_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    """
+    Dependency to extract and validate the API key from the Authorization header.
+    Expects the header in the format: Authorization: Bearer YOUR_API_KEY
+    """
+    if not credentials:
+        logger.warning("Authorization header missing")
+        raise HTTPException(
+            status_code=401,
+            detail="Authorization header missing",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    if credentials.scheme.lower() != "bearer":
+        logger.warning(f"Invalid authentication scheme: {credentials.scheme}")
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid authentication scheme. Expected 'Bearer'.",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    api_key = credentials.credentials
     if not api_key:
-        logger.warning("API key missing")
+        logger.warning("API key missing in Authorization header")
         raise HTTPException(
             status_code=401,
             detail="API key missing",
-            headers={"WWW-Authenticate": "API Key"},
+            headers={"WWW-Authenticate": "Bearer"},
         )
     if api_key not in API_KEYS:
         logger.warning(f"Invalid API key: {api_key}")
         raise HTTPException(
             status_code=403,
             detail="Invalid API key",
-            headers={"WWW-Authenticate": "API Key"},
+            headers={"WWW-Authenticate": "Bearer"},
         )
     return api_key