whoami-fastapi.py

#!/usr/bin/env python3
"""
whoami service authentication with the Hub
"""
import json
import os
from datetime import datetime
from typing import Optional
from urllib.parse import quote

from pydantic import BaseModel
from fastapi import Cookie, Header, Depends, FastAPI
from fastapi.responses import RedirectResponse

from jupyterhub.services.auth import HubAuth


prefix = os.environ.get('JUPYTERHUB_SERVICE_PREFIX', '/')

auth = HubAuth(api_token=os.environ['JUPYTERHUB_API_TOKEN'], cache_max_age=60)

app = FastAPI()


class RequiresLoginException(Exception):
    pass


class User(BaseModel):
    admin: bool
    last_activity: datetime
    name: str
    server: str
    pending: Optional[str] = None


@app.exception_handler(RequiresLoginException)
async def exception_handler(request: Request, exc: RequiresLoginException) -> Response:
    return RedirectResponse(url=auth.login_url + '?next=%s' % quote(request.path))


def header_or_cookie_user(
    authorization: Optional[str] = Header(None),
    cookie: Optional[str] = Cookie(None)
) -> User:
    if cookie:
        return User(auth.user_for_cookie(cookie))
    if authorization:
        return User(auth.user_for_token(authorization))

    raise RequiresLoginException()


@app.get(prefix)
def whoami(user: User = Depends(header_or_cookie_user)):
    return user