Spaces:
Sleeping
Sleeping
from fastapi import APIRouter, Depends, HTTPException, status | |
from fastapi.security import OAuth2PasswordRequestForm | |
from auth import authenticate_user, get_current_user | |
from custom_auth import ( | |
create_token_for_user, | |
get_current_user_from_token, | |
get_token, | |
remove_token, | |
) | |
from models import Token, UserCreate, User | |
from database import get_users, create_account | |
router = APIRouter() | |
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()): | |
""" | |
Endpoint to get an access token using username and password | |
""" | |
user = authenticate_user(get_users(), form_data.username, form_data.password) | |
if not user: | |
raise HTTPException( | |
status_code=status.HTTP_401_UNAUTHORIZED, | |
detail="Incorrect username or password", | |
headers={"WWW-Authenticate": "Bearer"}, | |
) | |
# Create a new token for the user (this will remove any existing token) | |
access_token = create_token_for_user(user.username) | |
return Token(access_token=access_token, token_type="bearer") | |
async def register_user(user_data: UserCreate): | |
""" | |
Endpoint to register a new account | |
""" | |
success, message = create_account(user_data.username, user_data.password) | |
if not success: | |
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=message) | |
return {"message": message} | |
async def logout( | |
current_user: User = Depends(get_current_user_from_token), | |
token: str = Depends(get_token), | |
): | |
""" | |
Endpoint to logout (invalidate the current token) | |
""" | |
success = remove_token(token) | |
if success: | |
return {"message": "Logout successful"} | |
else: | |
return {"message": "Token already invalid or expired"} | |
async def get_current_user_info( | |
current_user: User = Depends(get_current_user_from_token), | |
): | |
""" | |
Get the current user's information | |
""" | |
return { | |
"username": current_user.username, | |
"email": current_user.email, | |
"full_name": current_user.full_name, | |
} | |