Merge pull request #141 from bento-platform/feat/auth/more-helpers

feat(auth): add evaluate-to-dict auth middleware functions

bento_lib/auth/middleware/base.py

@@ -6,7 +6,7 @@
 from typing import Any, Callable, Iterable
 
 from ..exceptions import BentoAuthException
-from ..types import EvaluationResultMatrix
+from ..types import EvaluationResultMatrix, EvaluationResultDict
 from .mark_authz_done_mixin import MarkAuthzDoneMixin
 
 __all__ = ["BaseAuthMiddleware"]
@@ -98,9 +98,13 @@ def _evaluate_body(resources: Iterable[dict], permissions: Iterable[str]) -> dic
         return {"resources": tuple(resources), "permissions": tuple(permissions)}
 
     @staticmethod
-    def _matrix_tuple_cast(authz_result: list[list[bool]]) -> tuple[tuple[bool, ...], ...]:
+    def _matrix_tuple_cast(authz_result: list[list[bool]]) -> EvaluationResultMatrix:
         return tuple(tuple(x) for x in authz_result)
 
+    @staticmethod
+    def _permissions_matrix_to_dict(m: EvaluationResultMatrix, permissions: Iterable[str]) -> EvaluationResultDict:
+        return tuple({p: pv for p, pv in zip(permissions, r)} for r in m)
+
     def evaluate(
         self,
         request: Any,
@@ -122,6 +126,20 @@ def evaluate(
             )["result"]
         )
 
+    def evaluate_to_dict(
+        self,
+        request: Any,
+        resources: Iterable[dict],
+        permissions: Iterable[str],
+        require_token: bool = False,
+        headers_getter: Callable[[Any], dict[str, str]] | None = None,
+        mark_authz_done: bool = False,
+    ) -> EvaluationResultDict:
+        # consume iterable only once in case it's a generator
+        _perms = tuple(permissions)
+        return self._permissions_matrix_to_dict(
+            self.evaluate(request, resources, _perms, require_token, headers_getter, mark_authz_done), _perms)
+
     def evaluate_one(
         self,
         request: Any,
@@ -176,6 +194,21 @@ async def async_evaluate(
             )["result"]
         )
 
+    async def async_evaluate_to_dict(
+        self,
+        request: Any,
+        resources: Iterable[dict],
+        permissions: Iterable[str],
+        require_token: bool = False,
+        headers_getter: Callable[[Any], dict[str, str]] | None = None,
+        mark_authz_done: bool = False,
+    ) -> EvaluationResultDict:
+        # consume iterable only once in case it's a generator
+        _perms = tuple(permissions)
+        return self._permissions_matrix_to_dict(
+            await self.async_evaluate(request, resources, _perms, require_token, headers_getter, mark_authz_done),
+            _perms)
+
     async def async_evaluate_one(
         self,
         request: Any,

bento_lib/auth/types.py

@@ -1,13 +1,16 @@
 from typing import Type
 
 from ..auth.middleware.mark_authz_done_mixin import MarkAuthzDoneMixin
+from ..auth.permissions import Permission
 
 __all__ = [
     "EvaluationResultMatrix",
+    "EvaluationResultDict",
     "MarkAuthzDoneType",
 ]
 
 EvaluationResultMatrix = tuple[tuple[bool, ...], ...]
+EvaluationResultDict = tuple[dict[Permission, bool], ...]
 
 # Allow subclass OR instance, since mark_authz_done is a static method:
 MarkAuthzDoneType = MarkAuthzDoneMixin | Type[MarkAuthzDoneMixin]

tests/test_platform_fastapi.py

@@ -170,6 +170,20 @@ async def auth_post_with_token_evaluate_one(request: Request, body: TestTokenBod
     )})
 
 
+@test_app_auth.post("/post-with-token-evaluate-to-dict")
+async def auth_post_with_token_evaluate_to_dict(request: Request, body: TestTokenBody):
+    token = body.token
+
+    auth_middleware.mark_authz_done(request)
+    return JSONResponse({"payload": await auth_middleware.async_evaluate_to_dict(
+        request,
+        (RESOURCE_EVERYTHING,),
+        (PERMISSION_INGEST_DATA,),
+        require_token=True,
+        headers_getter=(lambda _r: {"Authorization": f"Bearer {token}"}),
+    )})
+
+
 # Auth test app (disabled auth middleware) ------------------------------------
 
 test_app_auth_disabled = FastAPI()
@@ -322,6 +336,13 @@ def test_fastapi_auth_post_with_token_evaluate_one(aioresponse: aioresponses, fa
     assert r.text == '{"payload":true}'
 
 
+def test_fastapi_auth_post_with_token_evaluate_to_dict(aioresponse: aioresponses, fastapi_client_auth: TestClient):
+    aioresponse.post("https://bento-auth.local/policy/evaluate", status=200, payload={"result": [[True]]})
+    r = fastapi_client_auth.post("/post-with-token-evaluate-to-dict", json={"token": "test"})
+    assert r.status_code == 200
+    assert r.text == '{"payload":[{"ingest:data":true}]}'
+
+
 @pytest.mark.asyncio
 async def test_fastapi_auth_disabled(aioresponse: aioresponses, fastapi_client_auth_disabled: TestClient):
     # middleware is disabled, should work anyway

tests/test_platform_flask.py

@@ -127,6 +127,19 @@ def auth_post_with_token_evaluate_one():
             headers_getter=(lambda _r: {"Authorization": f"Bearer {token}"}),
         )})
 
+    @test_app_auth.route("/post-with-token-evaluate-to-dict", methods=["POST"])
+    def auth_post_with_token_evaluate__to_dict():
+        token = request.json["token"]
+
+        auth_middleware.mark_authz_done(request)
+        return jsonify({"payload": auth_middleware.evaluate_to_dict(
+            request,
+            (RESOURCE_EVERYTHING,),
+            (PERMISSION_INGEST_DATA,),
+            require_token=True,
+            headers_getter=(lambda _r: {"Authorization": f"Bearer {token}"}),
+        )})
+
     with test_app_auth.test_client() as client:
         yield client
 
@@ -258,6 +271,14 @@ def test_flask_auth_post_with_token_evaluate_one(flask_client_auth: FlaskClient)
     assert r.text == '{"payload":true}\n'
 
 
+@responses.activate
+def test_flask_auth_post_with_token_evaluate_to_dict(flask_client_auth: FlaskClient):
+    responses.add(responses.POST, "https://bento-auth.local/policy/evaluate", json={"result": [[True]]}, status=200)
+    r = flask_client_auth.post("/post-with-token-evaluate-to-dict", json={"token": "test"})
+    assert r.status_code == 200
+    assert r.text == '{"payload":[{"ingest:data":true}]}\n'
+
+
 @responses.activate
 def test_flask_auth_disabled(flask_client_auth_disabled_with_middleware: tuple[FlaskClient, FlaskAuthMiddleware]):
     flask_client_auth_disabled, auth_middleware_disabled = flask_client_auth_disabled_with_middleware