File size: 2,145 Bytes
f3fa932 344b3b7 f3fa932 90cbfa2 f3fa932 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """
Auth Middleware — JWT Validation
"""
import os
import logging
from typing import Optional
from datetime import datetime, timedelta, timezone
from fastapi import Request, HTTPException, Depends
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from jose import JWTError, jwt
from bson import ObjectId
from src.db import mongo
logger = logging.getLogger("llm_firewall.auth_middleware")
security = HTTPBearer()
SECRET_KEY = os.getenv("JWT_SECRET_KEY", "d9f933e574ce4f4f41da2de40e9b3f2e72452ee35e906d8639f36da8033bb982")
ALGORITHM = os.getenv("JWT_ALGORITHM", "HS256")
ACCESS_TOKEN_EXPIRE_MINUTES = int(os.getenv("ACCESS_TOKEN_EXPIRE_MINUTES", "60"))
def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
to_encode = data.copy()
if expires_delta:
expire = datetime.now(timezone.utc) + expires_delta
else:
expire = datetime.now(timezone.utc) + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
to_encode.update({"exp": expire})
encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
return encoded_jwt
async def validate_user_token(credentials: HTTPAuthorizationCredentials = Depends(security)) -> dict:
"""
Validate the JWT access token and return the user document.
Raises HTTPException on failure.
"""
token = credentials.credentials
try:
payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
user_id_str: str = payload.get("sub")
if user_id_str is None:
raise HTTPException(status_code=401, detail="Invalid token payload")
try:
user_id = ObjectId(user_id_str)
except Exception:
raise HTTPException(status_code=401, detail="Invalid user ID format in token")
except JWTError:
raise HTTPException(status_code=401, detail="Could not validate credentials")
users_collection = mongo.get_users_collection()
user = await users_collection.find_one({"_id": user_id})
if user is None:
raise HTTPException(status_code=401, detail="User not found")
return user
|