| from fastapi import Depends, HTTPException, status | |
| from jose import jwt | |
| from .Model import User | |
| from .Constants import UserType | |
| from fastapi.security import OAuth2PasswordBearer | |
| SECRET_KEY = "your_secret_key_here" | |
| ALGORITHM = "HS256" | |
| oauth2_scheme = OAuth2PasswordBearer(tokenUrl="user/login") | |
| async def get_current_user(token: str = Depends(oauth2_scheme)): | |
| credentials_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]) | |
| user = await User.get_or_none(phoneNumber=payload.get("sub")) | |
| if user is None: | |
| raise credentials_exception | |
| except jwt.JWTError: | |
| raise credentials_exception | |
| return user | |
| async def get_current_active_user(current_user: User = Depends(get_current_user)): | |
| return current_user | |
| async def get_admin_user(current_user: User = Depends(get_current_active_user)): | |
| print(current_user.user_type, current_user.name) | |
| if current_user.user_type != UserType.ADMIN: | |
| raise HTTPException( | |
| status_code=status.HTTP_401_UNAUTHORIZED, | |
| detail="User does not have permission", | |
| headers={"WWW-Authenticate": "Bearer"}, | |
| ) | |
| return current_user | |