Spaces:
Sleeping
Sleeping
File size: 2,612 Bytes
8de5584 | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | 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
|