diff --git a/pingpong/authz/mock.py b/pingpong/authz/mock.py index a863666e8..1a87b40c6 100644 --- a/pingpong/authz/mock.py +++ b/pingpong/authz/mock.py @@ -53,6 +53,9 @@ def __init__(self, driver: OpenFgaAuthzDriver, params: dict | None = None): f"/stores/{self._test_store_id}/authorization-models/{self._test_model_id}" )(self._api_test_store_get_model) self.app.post(f"/stores/{self._test_store_id}/check")(self._api_check) + self.app.post(f"/stores/{self._test_store_id}/list-objects")( + self._api_list_objects + ) self.app.post(f"/stores/{self._test_store_id}/write")(self._api_write) self.app.get("/inspect/calls")(self._api_inspect_calls) @@ -98,6 +101,27 @@ async def _api_check(self, request: Request): "allowed": self._has_grant((user, relation, obj)), } + async def _api_list_objects(self, request: Request): + body = await request.json() + user = body.get("user") + relation = body.get("relation") + obj_type = body.get("type") + + if not user or not relation or not obj_type: + raise ValueError("Missing user, relation or type") + + prefix = f"{obj_type}:" + objects = [ + obj + for (u, rel, obj) in self._all_grants + if u == user and rel == relation and obj.startswith(prefix) + ] + + return { + "objects": objects, + "continuation_token": "", + } + async def _api_write(self, request: Request): body = await request.json() # Process added permissions diff --git a/pingpong/permission.py b/pingpong/permission.py index 0e0b1cf14..3859f37d9 100644 --- a/pingpong/permission.py +++ b/pingpong/permission.py @@ -171,3 +171,21 @@ async def test(self, request: Request) -> bool: def __str__(self): return f"Authz({self.relation}, {self.target})" + + +class InstitutionAdmin(Expression): + async def test(self, request: Request) -> bool: + if not hasattr(request.state, "auth_user") or not request.state.auth_user: + return False + + try: + institutions = await request.state.authz.list( + request.state.auth_user, "admin", "institution" + ) + return len(institutions) > 0 + except Exception as e: + logger.exception("Error evaluating expression %s: %s", self, e) + raise HTTPException(status_code=500, detail=str(e)) + + def __str__(self): + return "InstitutionAdmin()" diff --git a/pingpong/server.py b/pingpong/server.py index 1e7f4270a..feca98ecb 100644 --- a/pingpong/server.py +++ b/pingpong/server.py @@ -128,7 +128,7 @@ handle_delete_files, ) from .now import NowFn, utcnow -from .permission import Authz, LoggedIn +from .permission import Authz, InstitutionAdmin, LoggedIn from .runs import get_placeholder_ci_calls from .vector_stores import ( add_vector_store_files_to_db, @@ -869,7 +869,7 @@ async def auth(request: Request): @v1.get( "/api_keys/default", - dependencies=[Depends(Authz("admin"))], + dependencies=[Depends(Authz("admin") | InstitutionAdmin())], response_model=schemas.DefaultAPIKeys, ) async def list_default_api_keys(request: Request): diff --git a/pingpong/test_server.py b/pingpong/test_server.py index 3cbff1e79..9fcf393e4 100644 --- a/pingpong/test_server.py +++ b/pingpong/test_server.py @@ -171,6 +171,45 @@ async def test_config_correct_permissions(api, valid_user_token): assert response.status_code == 200 +@with_user(123) +@with_authz(grants=[]) +async def test_default_api_keys_requires_permissions(api, valid_user_token): + response = api.get( + "/api/v1/api_keys/default", + cookies={ + "session": valid_user_token, + }, + ) + assert response.status_code == 403 + assert response.json() == {"detail": "Missing required role"} + + +@with_user(123) +@with_authz(grants=[("user:123", "admin", "institution:1")]) +async def test_default_api_keys_allows_institution_admin(api, valid_user_token): + response = api.get( + "/api/v1/api_keys/default", + cookies={ + "session": valid_user_token, + }, + ) + assert response.status_code == 200 + assert response.json() == {"default_keys": []} + + +@with_user(123) +@with_authz(grants=[("user:123", "admin", "root:0")]) +async def test_default_api_keys_allows_root_admin(api, valid_user_token): + response = api.get( + "/api/v1/api_keys/default", + cookies={ + "session": valid_user_token, + }, + ) + assert response.status_code == 200 + assert response.json() == {"default_keys": []} + + async def test_auth_with_invalid_token(api): invalid_token = ( "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9." diff --git a/pingpong/testutil.py b/pingpong/testutil.py index dccaf6c5f..b9a6b6c8f 100644 --- a/pingpong/testutil.py +++ b/pingpong/testutil.py @@ -34,4 +34,6 @@ def with_authz_series(series): def with_authz(grants=None): - return with_authz_series([{"grants": grants}] if grants else None) + if grants is None: + return with_authz_series(None) + return with_authz_series([{"grants": grants}])