from datetime import datetime, timedelta from typing import Optional from fastapi import Depends, HTTPException, status from fastapi.security import OAuth2PasswordBearer import jwt from passlib.context import CryptContext from pydantic import BaseModel import os # Setup SECRET_KEY = os.getenv("SECRET_KEY", "change-this-key") ALGORITHM = os.getenv("ALGORITHM", "HS256") # Use pbkdf2 as fallback if bcrypt has issues on Windows pwd_context = CryptContext(schemes=["pbkdf2_sha256", "bcrypt"], deprecated="auto") oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token") # Models class Token(BaseModel): access_token: str token_type: str class TokenData(BaseModel): username: Optional[str] = None role: Optional[str] = None # Functions def verify_password(plain_password, hashed_password): """Verify a plain password against a hashed password""" return pwd_context.verify(plain_password, hashed_password) def get_password_hash(password): """Hash a password using bcrypt""" return pwd_context.hash(password) def create_access_token(data: dict, expires_delta: Optional[timedelta] = None): """Create a JWT access token""" to_encode = data.copy() if expires_delta: expire = datetime.utcnow() + expires_delta else: expire = datetime.utcnow() + timedelta(minutes=30) to_encode.update({"exp": expire}) encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM) return encoded_jwt async def get_current_user(token: str = Depends(oauth2_scheme)): """Extract and validate JWT token""" credential_exception = HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Could not validate credentials", headers={"WWW-Authenticate": "Bearer"}, ) try: payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM]) username: str = payload.get("sub") role: str = payload.get("role") if username is None: raise credential_exception token_data = TokenData(username=username, role=role) except jwt.InvalidTokenError: raise credential_exception return token_data def require_role(*allowed_roles): """Decorator to require specific user roles""" async def role_checker(current_user: TokenData = Depends(get_current_user)): if current_user.role not in allowed_roles: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=f"Insufficient permissions. Required roles: {allowed_roles}", ) return current_user return role_checker