From 51cc14140f117f88b78a6ba1a6a5ed64dc54184b Mon Sep 17 00:00:00 2001 From: Russ Allbery Date: Wed, 10 May 2023 15:11:28 -0700 Subject: [PATCH] Split out Gafaelfawr storage layer and use models Rather than embedding the Gafaelfawr client in the AuthenticatedUser model, separate out a Gafaelfawr storage client. Use copies of the Gafaelfawr models rather than sending and receiving raw JSON. These can be replaced with the real Gafaelfawr models once those are available from PyPI. Take advantage of this to use Pydantic for the parsing and capture Gafaelfawr parsing errors in a Slack-reportable exception. --- src/mobu/constants.py | 27 ++++ src/mobu/exceptions.py | 56 +++++++ src/mobu/factory.py | 12 +- src/mobu/models/user.py | 56 +------ src/mobu/services/flock.py | 7 +- src/mobu/services/manager.py | 12 +- src/mobu/services/solitary.py | 11 +- src/mobu/storage/gafaelfawr.py | 150 ++++++++++++++++++ tests/storage/__init__.py | 0 .../gafaelfawr_test.py} | 12 +- tests/support/gafaelfawr.py | 9 +- 11 files changed, 287 insertions(+), 65 deletions(-) create mode 100644 src/mobu/storage/gafaelfawr.py create mode 100644 tests/storage/__init__.py rename tests/{user_test.py => storage/gafaelfawr_test.py} (63%) diff --git a/src/mobu/constants.py b/src/mobu/constants.py index 41254dd3..7b610fec 100644 --- a/src/mobu/constants.py +++ b/src/mobu/constants.py @@ -1,7 +1,34 @@ """Global constants for mobu.""" +from __future__ import annotations + +from datetime import timedelta + +__all__ = [ + "NOTEBOOK_REPO_URL", + "NOTEBOOK_REPO_BRANCH", + "TOKEN_LIFETIME", + "USERNAME_REGEX", +] + NOTEBOOK_REPO_URL = "https://github.com/lsst-sqre/notebook-demo.git" """Default notebook repository for NotebookRunner.""" NOTEBOOK_REPO_BRANCH = "prod" """Default repository branch for NotebookRunner.""" + +TOKEN_LIFETIME = timedelta(days=365) +"""Token lifetime for mobu's service tokens. + +mobu currently has no mechanism for refreshing tokens while running, so this +should be long enough that mobu will be restarted before the tokens expire. +An expiration exists primarily to ensure that the tokens don't accumulate +forever. +""" + +# This must be kept in sync with Gafaelfawr until we can import the models +# from Gafaelfawr directly. +USERNAME_REGEX = ( + "^[a-z0-9](?:[a-z0-9]|-[a-z0-9])*[a-z](?:[a-z0-9]|-[a-z0-9])*$" +) +"""Regex matching all valid usernames.""" diff --git a/src/mobu/exceptions.py b/src/mobu/exceptions.py index 8731238c..787b3196 100644 --- a/src/mobu/exceptions.py +++ b/src/mobu/exceptions.py @@ -8,6 +8,7 @@ from fastapi import status from httpx_ws import HTTPXWSException, WebSocketDisconnect +from pydantic import ValidationError from safir.datetime import format_datetime_for_logging from safir.fastapi import ClientRequestError from safir.models import ErrorLocation @@ -29,6 +30,7 @@ "CachemachineError", "CodeExecutionError", "FlockNotFoundError", + "GafaelfawrParseError", "GafaelfawrWebError", "JupyterTimeoutError", "JupyterWebError", @@ -62,6 +64,60 @@ def _remove_ansi_escapes(string: str) -> str: return _ANSI_REGEX.sub("", string) +class GafaelfawrParseError(SlackException): + """Unable to parse the reply from Gafaelfawr. + + Parameters + ---------- + message + Summary error message. + error + Detailed error message, possibly multi-line. + user + Username of the user involved. + """ + + @classmethod + def from_exception( + cls, exc: ValidationError, user: Optional[str] = None + ) -> Self: + """Create an exception from a Pydantic parse failure. + + Parameters + ---------- + exc + Pydantic exception. + user + Username of the user involved. + + Returns + ------- + GafaelfawrParseError + Constructed exception. + """ + error = f"{type(exc).__name__}: {str(exc)}" + return cls("Unable to parse reply from Gafalefawr", error, user) + + def __init__( + self, message: str, error: str, user: Optional[str] = None + ) -> None: + super().__init__(message, user) + self.error = error + + def to_slack(self) -> SlackMessage: + """Convert to a Slack message for Slack alerting. + + Returns + ------- + SlackMessage + Slack message suitable for posting as an alert. + """ + message = super().to_slack() + block = SlackCodeBlock(heading="Error", code=self.error) + message.blocks.append(block) + return message + + class GafaelfawrWebError(SlackWebException): """An API call to Gafaelfawr failed.""" diff --git a/src/mobu/factory.py b/src/mobu/factory.py index 5a1d85f8..e3c7aaf5 100644 --- a/src/mobu/factory.py +++ b/src/mobu/factory.py @@ -13,6 +13,7 @@ from .models.solitary import SolitaryConfig from .services.manager import FlockManager from .services.solitary import Solitary +from .storage.gafaelfawr import GafaelfawrStorage __all__ = ["Factory", "ProcessContext"] @@ -38,7 +39,9 @@ class ProcessContext: def __init__(self, http_client: AsyncClient) -> None: self.http_client = http_client - self.manager = FlockManager(http_client, structlog.get_logger("mobu")) + logger = structlog.get_logger("mobu") + gafaelfawr = GafaelfawrStorage(http_client, logger) + self.manager = FlockManager(gafaelfawr, http_client, logger) async def aclose(self) -> None: """Clean up a process context. @@ -93,5 +96,10 @@ def create_solitary(self, solitary_config: SolitaryConfig) -> Solitary: Newly-created solitary manager. """ return Solitary( - solitary_config, self._context.http_client, self._logger + solitary_config=solitary_config, + gafaelfawr_storage=GafaelfawrStorage( + self._context.http_client, self._logger + ), + http_client=self._context.http_client, + logger=self._logger, ) diff --git a/src/mobu/models/user.py b/src/mobu/models/user.py index 5c22e96d..73b81979 100644 --- a/src/mobu/models/user.py +++ b/src/mobu/models/user.py @@ -2,16 +2,15 @@ from __future__ import annotations -import time -from typing import Any, Optional, Self +from typing import Optional -from httpx import AsyncClient, HTTPError from pydantic import BaseModel, Field -from ..config import config -from ..exceptions import GafaelfawrWebError - -__all__ = ["AuthenticatedUser", "User", "UserSpec"] +__all__ = [ + "AuthenticatedUser", + "User", + "UserSpec", +] class User(BaseModel): @@ -91,46 +90,3 @@ class AuthenticatedUser(User): title="Authentication token for user", example="gt-1PhgAeB-9Fsa-N1NhuTu_w.oRvMvAQp1bWfx8KCJKNohg", ) - - @classmethod - async def create( - cls, user: User, scopes: list[str], http_client: AsyncClient - ) -> Self: - if not config.environment_url: - raise RuntimeError("environment_url not set") - token_url = ( - str(config.environment_url).rstrip("/") + "/auth/api/v1/tokens" - ) - data: dict[str, Any] = { - "username": user.username, - "name": "Mobu Test User", - "token_type": "user", - "token_name": f"mobu {str(float(time.time()))}", - "scopes": scopes, - "expires": int(time.time() + 60 * 60 * 24 * 365), - } - if user.uidnumber is not None: - data["uid"] = user.uidnumber - if user.gidnumber is not None: - data["gid"] = user.gidnumber - else: - data["gid"] = user.uidnumber - elif user.gidnumber is not None: - data["gid"] = user.gidnumber - try: - r = await http_client.post( - token_url, - headers={"Authorization": f"Bearer {config.gafaelfawr_token}"}, - json=data, - ) - r.raise_for_status() - except HTTPError as e: - raise GafaelfawrWebError.from_exception(e, user.username) - body = r.json() - return cls( - username=user.username, - uidnumber=data["uid"] if "uid" in data else None, - gidnumber=data["gid"] if "gid" in data else None, - token=body["token"], - scopes=scopes, - ) diff --git a/src/mobu/services/flock.py b/src/mobu/services/flock.py index 7f04f233..70a3693b 100644 --- a/src/mobu/services/flock.py +++ b/src/mobu/services/flock.py @@ -15,6 +15,7 @@ from ..exceptions import MonkeyNotFoundError from ..models.flock import FlockConfig, FlockData, FlockSummary from ..models.user import AuthenticatedUser, User, UserSpec +from ..storage.gafaelfawr import GafaelfawrStorage from .monkey import Monkey __all__ = ["Flock"] @@ -29,6 +30,8 @@ class Flock: Configuration for this flock of monkeys. scheduler Job scheduler used to manage the tasks for the monkeys. + gafaelfawr_storage + Gafaelfawr storage client. http_client Shared HTTP client. logger @@ -40,12 +43,14 @@ def __init__( *, flock_config: FlockConfig, scheduler: Scheduler, + gafaelfawr_storage: GafaelfawrStorage, http_client: AsyncClient, logger: BoundLogger, ) -> None: self.name = flock_config.name self._config = flock_config self._scheduler = scheduler + self._gafaelfawr = gafaelfawr_storage self._http_client = http_client self._logger = logger.bind(flock=self.name) self._monkeys: dict[str, Monkey] = {} @@ -146,7 +151,7 @@ async def _create_users(self) -> list[AuthenticatedUser]: users = self._users_from_spec(self._config.user_spec, count) scopes = self._config.scopes return [ - await AuthenticatedUser.create(u, scopes, self._http_client) + await self._gafaelfawr.create_service_token(u, scopes) for u in users ] diff --git a/src/mobu/services/manager.py b/src/mobu/services/manager.py index 7d722ef4..da876b23 100644 --- a/src/mobu/services/manager.py +++ b/src/mobu/services/manager.py @@ -12,6 +12,7 @@ from ..config import config from ..exceptions import FlockNotFoundError from ..models.flock import FlockConfig, FlockSummary +from ..storage.gafaelfawr import GafaelfawrStorage from .flock import Flock __all__ = ["FlockManager"] @@ -26,13 +27,21 @@ class FlockManager: Parameters ---------- + gafaelfawr_storage + Gafaelfawr storage client. http_client Shared HTTP client. logger Global logger to use for process-wide (not monkey) logging. """ - def __init__(self, http_client: AsyncClient, logger: BoundLogger) -> None: + def __init__( + self, + gafaelfawr_storage: GafaelfawrStorage, + http_client: AsyncClient, + logger: BoundLogger, + ) -> None: + self._gafaelfawr = gafaelfawr_storage self._http_client = http_client self._logger = logger self._flocks: dict[str, Flock] = {} @@ -74,6 +83,7 @@ async def start_flock(self, flock_config: FlockConfig) -> Flock: flock = Flock( flock_config=flock_config, scheduler=self._scheduler, + gafaelfawr_storage=self._gafaelfawr, http_client=self._http_client, logger=self._logger, ) diff --git a/src/mobu/services/solitary.py b/src/mobu/services/solitary.py index b00a6444..c5362aa3 100644 --- a/src/mobu/services/solitary.py +++ b/src/mobu/services/solitary.py @@ -8,7 +8,7 @@ from structlog.stdlib import BoundLogger from ..models.solitary import SolitaryConfig, SolitaryResult -from ..models.user import AuthenticatedUser +from ..storage.gafaelfawr import GafaelfawrStorage from .monkey import Monkey __all__ = ["Solitary"] @@ -21,6 +21,8 @@ class Solitary: ---------- solitary_config Configuration for the monkey. + gafaelfawr_storage + Gafaelfawr storage client. http_client Shared HTTP client. logger @@ -29,11 +31,14 @@ class Solitary: def __init__( self, + *, solitary_config: SolitaryConfig, + gafaelfawr_storage: GafaelfawrStorage, http_client: AsyncClient, logger: BoundLogger, ) -> None: self._config = solitary_config + self._gafaelfawr = gafaelfawr_storage self._http_client = http_client self._logger = logger @@ -45,8 +50,8 @@ async def run(self) -> SolitaryResult: SolitaryResult Result of monkey run. """ - user = await AuthenticatedUser.create( - self._config.user, self._config.scopes, self._http_client + user = await self._gafaelfawr.create_service_token( + self._config.user, self._config.scopes ) monkey = Monkey( name=f"solitary-{user.username}", diff --git a/src/mobu/storage/gafaelfawr.py b/src/mobu/storage/gafaelfawr.py new file mode 100644 index 00000000..149a66f2 --- /dev/null +++ b/src/mobu/storage/gafaelfawr.py @@ -0,0 +1,150 @@ +"""Manage Gafaelfawr users and tokens.""" + +from __future__ import annotations + +import json +from datetime import datetime +from enum import Enum +from typing import Optional + +from httpx import AsyncClient, HTTPError +from pydantic import BaseModel, Field, ValidationError +from safir.datetime import current_datetime +from structlog.stdlib import BoundLogger + +from ..config import config +from ..constants import TOKEN_LIFETIME, USERNAME_REGEX +from ..exceptions import GafaelfawrParseError, GafaelfawrWebError +from ..models.user import AuthenticatedUser, User + +__all__ = ["GafaelfawrStorage"] + + +class _TokenType(Enum): + """The class of token. + + This is copied from Gafaelfawr and should be replaced with using the + Gafaelfawr models directly once they're available. + """ + + session = "session" + user = "user" + notebook = "notebook" + internal = "internal" + service = "service" + + +class _AdminTokenRequest(BaseModel): + """Request by a Gafaelfawr token administrator to create a token. + + This is copied from Gafaelfawr and should be replaced with using the + Gafaelfawr models directly once they're available. + """ + + username: str = Field( + ..., min_length=1, max_length=64, regex=USERNAME_REGEX + ) + token_type: _TokenType = Field(...) + scopes: list[str] = Field([]) + expires: Optional[datetime] = Field(None) + name: Optional[str] = Field(None, min_length=1) + uid: Optional[int] = Field(None, ge=1) + gid: Optional[int] = Field(None, ge=1) + + +class _NewToken(BaseModel): + """Response to a token creation request. + + This is copied from Gafaelfawr and should be replaced with using the + Gafaelfawr models directly once they're available. + """ + + token: str = Field(...) + + +class GafaelfawrStorage: + """Manage users and authentication tokens. + + mobu uses bot users to run its tests. Those users may be pre-existing or + manufactured on the fly by mobu. Either way, mobu creates new service + tokens for the configured users, and then provides those usernames and + tokens to monkeys to use for executing their business. + + This class handles the call to Gafaelfawr to create the service token. + + Parameters + ---------- + http_client + Shared HTTP client. + logger + Logger to use. + """ + + def __init__(self, http_client: AsyncClient, logger: BoundLogger) -> None: + self._client = http_client + self._logger = logger + + if not config.environment_url: + raise RuntimeError("environment_url not set") + base_url = str(config.environment_url).rstrip("/") + self._token_url = base_url + "/auth/api/v1/tokens" + + async def create_service_token( + self, user: User, scopes: list[str] + ) -> AuthenticatedUser: + """Create a service token for a user. + + Parameters + ---------- + user + Metadata for the user. If ``uid`` or ``gid`` are set for the user, + they will be stored with the token and override Gafaelfawr's + normal user metadata source. + scopes + Scopes the requested token should have. + + Returns + ------- + AuthenticatedUser + Authenticated user with their metadata, scopes, and token. + + Raises + ------ + GafaelfawrParseError + Raised if the input or output data for Gafaelfawr's token call + could not be parsed. + GafaelfawrWebError + Raised if an HTTP protocol error occurred talking to Gafaelfawr. + """ + request = _AdminTokenRequest( + username=user.username, + token_type=_TokenType.service, + scopes=scopes, + expires=current_datetime() + TOKEN_LIFETIME, + name="Mobu Test User", + uid=user.uidnumber, + gid=user.gidnumber or user.uidnumber, + ) + try: + # The awkward JSON generation ensures that datetime fields are + # properly serialized using the model rather than left as datetime + # objects, which httpx will not understand. Pydantic v2 will offer + # a better way to do this. + r = await self._client.post( + self._token_url, + headers={"Authorization": f"Bearer {config.gafaelfawr_token}"}, + json=json.loads(request.json(exclude_none=True)), + ) + r.raise_for_status() + token = _NewToken.parse_obj(r.json()) + return AuthenticatedUser( + username=user.username, + uidnumber=request.uid, + gidnumber=request.gid, + token=token.token, + scopes=scopes, + ) + except HTTPError as e: + raise GafaelfawrWebError.from_exception(e, user.username) + except ValidationError as e: + raise GafaelfawrParseError.from_exception(e, user.username) diff --git a/tests/storage/__init__.py b/tests/storage/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/user_test.py b/tests/storage/gafaelfawr_test.py similarity index 63% rename from tests/user_test.py rename to tests/storage/gafaelfawr_test.py index ee5fa13f..5095a116 100644 --- a/tests/user_test.py +++ b/tests/storage/gafaelfawr_test.py @@ -4,11 +4,13 @@ import pytest import respx +import structlog from safir.dependencies.http_client import http_client_dependency -from mobu.models.user import AuthenticatedUser, User +from mobu.models.user import User +from mobu.storage.gafaelfawr import GafaelfawrStorage -from .support.gafaelfawr import mock_gafaelfawr +from ..support.gafaelfawr import mock_gafaelfawr @pytest.mark.asyncio @@ -18,8 +20,12 @@ async def test_generate_token(respx_mock: respx.Router) -> None: scopes = ["exec:notebook"] client = await http_client_dependency() - user = await AuthenticatedUser.create(config, scopes, client) + logger = structlog.get_logger(__file__) + gafaelfawr = GafaelfawrStorage(client, logger) + + user = await gafaelfawr.create_service_token(config, scopes) assert user.username == "someuser" assert user.uidnumber == 1234 + assert user.gidnumber == 1234 assert user.scopes == ["exec:notebook"] assert user.token.startswith("gt-") diff --git a/tests/support/gafaelfawr.py b/tests/support/gafaelfawr.py index 46a7c86e..9a762a5a 100644 --- a/tests/support/gafaelfawr.py +++ b/tests/support/gafaelfawr.py @@ -5,12 +5,13 @@ import base64 import json import os -import time +from datetime import datetime from typing import Optional from unittest.mock import ANY import respx from httpx import Request, Response +from safir.datetime import current_datetime from mobu.config import config @@ -54,8 +55,7 @@ def handler(request: Request) -> Response: assert request.headers["Authorization"] == f"Bearer {admin_token}" expected = { "username": username if username else ANY, - "token_type": "user", - "token_name": ANY, + "token_type": "service", "scopes": ["exec:notebook"], "expires": ANY, "name": "Mobu Test User", @@ -69,8 +69,7 @@ def handler(request: Request) -> Response: expected["gid"] = ANY body = json.loads(request.content) assert body == expected - assert body["token_name"].startswith("mobu ") - assert body["expires"] > time.time() + assert datetime.fromisoformat(body["expires"]) > current_datetime() response = {"token": make_gafaelfawr_token(body["username"])} return Response(200, json=response)