Repository: aipotheosis-labs/aci
File: aipolabs/server/routes/auth.py (Python)

Code before:
    jwt_token: str = jwt.encode(
        header, payload, config.SIGNING_KEY
    ).decode()  # for this jwt lib, need to decode to convert bytes to string

    return jwt_token


# login route for different identity providers
@router.get("/login/{provider}", include_in_schema=True)
async def login(request: Request, provider: ClientIdentityProvider) -> RedirectResponse:
    oauth2_client = OAUTH2_CLIENTS[provider]

    path = request.url_for(LOGIN_CALLBACK_PATH_NAME, provider=provider.value).path
    redirect_uri = f"{config.AIPOLABS_REDIRECT_URI_BASE}{path}"
    logger.info(f"initiating login for provider={provider}, redirecting to={redirect_uri}")

    return await oauth2.authorize_redirect(oauth2_client, request, redirect_uri)


@router.get("/signup/{provider}", include_in_schema=True)
async def signup(
    request: Request,
    provider: ClientIdentityProvider,
    signup_code: str,
    db_session: Annotated[Session, Depends(deps.yield_db_session)],
) -> RedirectResponse:
    _validate_signup(db_session, signup_code)
    oauth2_client = OAUTH2_CLIENTS[provider]

    path = request.url_for(SIGNUP_CALLBACK_PATH_NAME, provider=provider.value).path
    redirect_uri = f"{config.AIPOLABS_REDIRECT_URI_BASE}{path}?signup_code={signup_code}"
    logger.info(
        f"initiating signup for provider={provider}, signup_code={signup_code}, redirecting to={redirect_uri}"
    )

    return await oauth2.authorize_redirect(oauth2_client, request, redirect_uri)


@router.get(
    "/signup/callback/{provider}",
    name=SIGNUP_CALLBACK_PATH_NAME,
    include_in_schema=True,
)
async def signup_callback(
    request: Request,
    response: Response,
    provider: ClientIdentityProvider,
    signup_code: str,
    db_session: Annotated[Session, Depends(deps.yield_db_session)],
) -> RedirectResponse:
    logger.info(

Review comment (suggestion):
are you creating pydantic class `SignUpCodeValidate` solely for making it a body parameter? If that's the case we should probably just use `signup_code: Annotated[str, Body()]`. 

Code after:

    # Authlib expects a header, payload, and key
    header = {"alg": config.JWT_ALGORITHM}
    # TODO: try/except, retry?
    jwt_token: str = jwt.encode(
        header, payload, config.SIGNING_KEY
    ).decode()  # for this jwt lib, need to decode to convert bytes to string

    return jwt_token


# login route for different identity providers
@router.get("/login/{provider}", include_in_schema=True)
async def login(request: Request, provider: ClientIdentityProvider) -> RedirectResponse:
    oauth2_client = OAUTH2_CLIENTS[provider]

    path = request.url_for(LOGIN_CALLBACK_PATH_NAME, provider=provider.value).path
    redirect_uri = f"{config.AIPOLABS_REDIRECT_URI_BASE}{path}"
    logger.info(f"initiating login for provider={provider}, redirecting to={redirect_uri}")

    return await oauth2.authorize_redirect(oauth2_client, request, redirect_uri)


@router.post("/validate-signup-code/", include_in_schema=True)
async def check_signup_code(
    body: SignUpCodeValidate,
    db_session: Annotated[Session, Depends(deps.yield_db_session)],
) -> Response:
    _validate_signup(db_session, body.signup_code)
    return Response(status_code=status.HTTP_200_OK)


@router.get("/signup/{provider}", include_in_schema=True)
async def signup(
    request: Request,
    provider: ClientIdentityProvider,
    signup_code: str,
    db_session: Annotated[Session, Depends(deps.yield_db_session)],
) -> RedirectResponse:
    _validate_signup(db_session, signup_code)
    oauth2_client = OAUTH2_CLIENTS[provider]

    request.session["signup_code"] = signup_code

    path = request.url_for(SIGNUP_CALLBACK_PATH_NAME, provider=provider.value).path
    redirect_uri = f"{config.AIPOLABS_REDIRECT_URI_BASE}{path}"
    logger.info(
        f"initiating signup for provider={provider}, signup_code={signup_code}, redirecting to={redirect_uri}"
    )

    return await oauth2.authorize_redirect(oauth2_client, request, redirect_uri)