auth.py

import os
import httpx

from authlib.integrations.starlette_client import OAuth
from fastapi import FastAPI
from fastapi.requests import Request
from fastapi.responses import RedirectResponse
from starlette.middleware.sessions import SessionMiddleware

OAUTH_CLIENT_ID = os.environ.get("OAUTH_CLIENT_ID")
OAUTH_CLIENT_SECRET = os.environ.get("OAUTH_CLIENT_SECRET")
OAUTH_SCOPES = os.environ.get("OAUTH_SCOPES")
OAUTH_SCOPES = "profile" # TODO: remove when openid is fixed (honor nonce)
OPENID_PROVIDER_URL =  os.environ.get("OPENID_PROVIDER_URL")

for value in (OAUTH_CLIENT_ID, OAUTH_CLIENT_SECRET, OAUTH_SCOPES, OPENID_PROVIDER_URL):
    if value is None:
        raise ValueError("Missing environment variable")

AUTHORIZE_URL = OPENID_PROVIDER_URL + "/oauth/authorize"
ACCESS_TOKEN_URL = OPENID_PROVIDER_URL + "/oauth/token"
USER_INFO_URL = OPENID_PROVIDER_URL + "/oauth/userinfo"

oauth = OAuth()
oauth.register(
    name="huggingface",
    client_id=OAUTH_CLIENT_ID,
    client_secret=OAUTH_CLIENT_SECRET,
    access_token_url=ACCESS_TOKEN_URL,
    authorize_url=AUTHORIZE_URL,
    api_base_url=OPENID_PROVIDER_URL,
    client_kwargs={"scope": OAUTH_SCOPES},
)


async def landing(request: Request):
    if request.session.get("user"):
        return RedirectResponse("/gradio")
    else:
        return RedirectResponse(request.url_for("oauth_login"))


async def oauth_login(request: Request):
    redirect_uri = request.url_for("oauth_redirect_callback")
    return await oauth.huggingface.authorize_redirect(request, redirect_uri)


async def oauth_redirect_callback(request: Request):
    token = await oauth.huggingface.authorize_access_token(request)

    async with httpx.AsyncClient() as client:
        resp = await client.get(USER_INFO_URL, headers={"Authorization": f"Bearer {token['access_token']}"})
        user_info = resp.json()

    request.session["user"] = user_info # TODO: we should store token instead
    return RedirectResponse(request.url_for("landing"))


async def check_oauth(request: Request, call_next):
    if request.url.path in (
        "/",
        "/auth/huggingface",
        "/auth/callback",
    ):  # not protected
        return await call_next(request)
    if request.session.get("user"):  # authenticated
        return await call_next(request)
    return RedirectResponse("/")


def get_app() -> FastAPI:
    app = FastAPI()
    app.middleware("http")(check_oauth)
    app.add_middleware(SessionMiddleware, secret_key="session-secret-key")  # TODO: make this is secret key
    app.get("/")(landing)
    app.get("/auth/huggingface")(oauth_login)
    app.get("/auth/callback")(oauth_redirect_callback)
    return app