kiwi-vpn/api/kiwi_vpn_api/routers/user.py

from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from pydantic import BaseModel
from sqlalchemy.orm import Session

from ..config import Config
from ..db import Connection, schemas

router = APIRouter(prefix="/user")
SCHEME = OAuth2PasswordBearer(
    tokenUrl=f".{router.prefix}/auth",
)


class Token(BaseModel):
    access_token: str
    token_type: str


@router.post("/auth", response_model=Token)
async def login(
    form_data: OAuth2PasswordRequestForm = Depends(),
    config: Config = Depends(Config.load),
    db: Session = Depends(Connection.get),
):
    user = schemas.User.login(
        db=db,
        name=form_data.username,
        password=form_data.password,
        crypt_context=await config.crypto.crypt_context,
    )

    if user is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Could not validate credentials",
            headers={"WWW-Authenticate": "Bearer"},
        )

    access_token = await config.jwt.create_token(user.name)
    return {"access_token": access_token, "token_type": "bearer"}


async def dep_get_current_user(
    token: str = Depends(SCHEME),
    db: Session = Depends(Connection.get),
    config: Config = Depends(Config.load),
):
    username = await config.jwt.decode_token(token)
    user = schemas.User.from_db(db, username)

    return user


@router.get("/current", response_model=schemas.User)
async def get_current_user(
    current_user: schemas.User = Depends(dep_get_current_user),
):
    if current_user is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Could not validate credentials",
            headers={"WWW-Authenticate": "Bearer"},
        )

    return current_user