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)