brewman/brewman/brewman/routers/login.py

from datetime import timedelta

from fastapi import (
    APIRouter,
    Cookie,
    Depends,
    Form,
    HTTPException,
    Response,
    Security,
    status,
)
from fastapi.responses import JSONResponse
from fastapi.security import OAuth2PasswordRequestForm
from sqlalchemy.orm import Session

from .. import __version__
from ..core.config import settings
from ..core.security import (
    Token,
    authenticate_user,
    client_allowed,
    create_access_token,
    get_current_active_user,
)
from ..db.session import SessionLocal
from ..schemas.auth import UserToken


router = APIRouter()


# Dependency
def get_db():
    try:
        db = SessionLocal()
        yield db
    finally:
        db.close()


@router.post("/token", response_model=Token)
async def login_for_access_token(
    response: Response,
    form_data: OAuth2PasswordRequestForm = Depends(),
    client_id: int = Cookie(None),
    otp: int = Form(None),
    db: Session = Depends(get_db),
):
    user = authenticate_user(form_data.username, form_data.password, client_id, otp, db)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Bearer"},
        )
    allowed, c_id = client_allowed(user, client_id, otp, db)
    db.commit()
    if c_id and c_id != client_id:
        response.set_cookie(
            key="client_id", value=str(c_id), max_age=10 * 365 * 24 * 60 * 60
        )
    if not allowed:
        not_allowed_response = JSONResponse(
            status_code=status.HTTP_401_UNAUTHORIZED,
            headers={"WWW-Authenticate": "Bearer"},
            content={"detail": "Client is not registered"},
        )
        not_allowed_response.set_cookie(
            key="client_id", value=str(c_id), max_age=10 * 365 * 24 * 60 * 60
        )
        return not_allowed_response
    access_token_expires = timedelta(minutes=settings.JWT_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={
            "sub": user.name,
            "scopes": ["authenticated"]
            + list(
                set(
                    [
                        p.name.replace(" ", "-").lower()
                        for r in user.roles
                        for p in r.permissions
                    ]
                )
            ),
            "userId": str(user.id),
            "lockedOut": user.locked_out,
            "ver": __version__.__version__,
        },
        expires_delta=access_token_expires,
    )
    return {"access_token": access_token, "token_type": "bearer"}


@router.post("/refresh", response_model=Token)
async def refresh_token(user: UserToken = Security(get_current_active_user)):
    access_token_expires = timedelta(minutes=settings.JWT_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={
            "sub": user.name,
            "scopes": user.permissions,
            "userId": str(user.id_),
            "lockedOut": user.locked_out,
            "ver": __version__.__version__,
        },
        expires_delta=access_token_expires,
    )
    return {"access_token": access_token, "token_type": "bearer"}