Implement menu authorization in KeycloakAuthManager (#51845)

Author: vincbeck

providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/commands.py

@@ -23,6 +23,7 @@
 from keycloak import KeycloakAdmin, KeycloakError
 
 from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod
+from airflow.api_fastapi.common.types import MenuItem
 from airflow.configuration import conf
 from airflow.providers.keycloak.auth_manager.constants import (
     CONF_CLIENT_ID_KEY,
@@ -70,7 +71,7 @@ def create_permissions_command(args):
 @cli_utils.action_cli
 @providers_configuration_loaded
 def create_all_command(args):
-    """Create Keycloak auth manager scopes in Keycloak."""
+    """Create all Keycloak auth manager entities in Keycloak."""
     client = _get_client(args)
     client_uuid = _get_client_uuid(args)
 
@@ -115,12 +116,11 @@ def _create_scopes(client: KeycloakAdmin, client_uuid: str):
 
 
 def _create_resources(client: KeycloakAdmin, client_uuid: str):
-    # Fetch existing scopes
     all_scopes = client.get_client_authz_scopes(client_uuid)
     scopes = [
         {"id": scope["id"], "name": scope["name"]}
         for scope in all_scopes
-        if scope["name"] in get_args(ResourceMethod)
+        if scope["name"] in ["GET", "POST", "PUT", "DELETE"]
     ]
 
     for resource in KeycloakResource:
@@ -133,6 +133,19 @@ def _create_resources(client: KeycloakAdmin, client_uuid: str):
             skip_exists=True,
         )
 
+    # Create menu item resources
+    scopes = [{"id": scope["id"], "name": scope["name"]} for scope in all_scopes if scope["name"] == "MENU"]
+
+    for item in MenuItem:
+        client.create_client_authz_resource(
+            client_id=client_uuid,
+            payload={
+                "name": item.value,
+                "scopes": scopes,
+            },
+            skip_exists=True,
+        )
+
     print("Resources created successfully.")
 
 

providers/keycloak/src/airflow/providers/keycloak/auth_manager/keycloak_auth_manager.py

@@ -18,14 +18,15 @@
 
 import json
 import logging
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 from urllib.parse import urljoin
 
 import requests
 from fastapi import FastAPI
 
 from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
 from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager
+from airflow.api_fastapi.common.types import MenuItem
 from airflow.cli.cli_config import CLICommand, GroupCommand
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
@@ -54,7 +55,6 @@
         PoolDetails,
         VariableDetails,
     )
-    from airflow.api_fastapi.common.types import MenuItem
 
 log = logging.getLogger(__name__)
 
@@ -198,7 +198,11 @@ def is_authorized_custom_view(
     def filter_authorized_menu_items(
         self, menu_items: list[MenuItem], *, user: KeycloakAuthManagerUser
     ) -> list[MenuItem]:
-        return menu_items
+        authorized_menus = self._is_batch_authorized(
+            permissions=[(cast("ResourceMethod", "MENU"), menu_item.value) for menu_item in menu_items],
+            user=user,
+        )
+        return [MenuItem(menu[1]) for menu in authorized_menus]
 
     def get_fastapi_app(self) -> FastAPI | None:
         from airflow.providers.keycloak.auth_manager.routes.login import login_router
@@ -260,6 +264,33 @@ def _is_authorized(
             )
         raise AirflowException(f"Unexpected error: {resp.status_code} - {resp.text}")
 
+    def _is_batch_authorized(
+        self,
+        *,
+        permissions: list[tuple[ResourceMethod, str]],
+        user: KeycloakAuthManagerUser,
+    ) -> set[tuple[ResourceMethod, str]]:
+        client_id = conf.get(CONF_SECTION_NAME, CONF_CLIENT_ID_KEY)
+        realm = conf.get(CONF_SECTION_NAME, CONF_REALM_KEY)
+        server_url = conf.get(CONF_SECTION_NAME, CONF_SERVER_URL_KEY)
+
+        resp = requests.post(
+            self._get_token_url(server_url, realm),
+            data=self._get_batch_payload(client_id, permissions),
+            headers=self._get_headers(user.access_token),
+        )
+
+        if resp.status_code == 200:
+            return {(perm["scopes"][0], perm["rsname"]) for perm in resp.json()}
+        if resp.status_code == 403:
+            return set()
+        if resp.status_code == 400:
+            error = json.loads(resp.text)
+            raise AirflowException(
+                f"Request not recognized by Keycloak. {error.get('error')}. {error.get('error_description')}"
+            )
+        raise AirflowException(f"Unexpected error: {resp.status_code} - {resp.text}")
+
     @staticmethod
     def _get_token_url(server_url, realm):
         return f"{server_url}/realms/{realm}/protocol/openid-connect/token"
@@ -276,6 +307,17 @@ def _get_payload(client_id: str, permission: str, attributes: dict[str, str] | N
 
         return payload
 
+    @staticmethod
+    def _get_batch_payload(client_id: str, permissions: list[tuple[ResourceMethod, str]]):
+        payload: dict[str, Any] = {
+            "grant_type": "urn:ietf:params:oauth:grant-type:uma-ticket",
+            "audience": client_id,
+            "permission": [f"{permission[1]}#{permission[0]}" for permission in permissions],
+            "response_mode": "permissions",
+        }
+
+        return payload
+
     @staticmethod
     def _get_headers(access_token):
         return {

providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_commands.py

@@ -23,6 +23,7 @@
 import pytest
 
 from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod
+from airflow.api_fastapi.common.types import MenuItem
 from airflow.cli import cli_parser
 from airflow.providers.keycloak.auth_manager.cli.commands import (
     _create_admin_permission,
@@ -116,7 +117,7 @@ def test_create_scopes_with_client_not_found(self, mock_get_client):
     def test_create_resources(self, mock_get_client):
         client = Mock()
         mock_get_client.return_value = client
-        scopes = [{"id": "1", "name": "GET"}]
+        scopes = [{"id": "1", "name": "GET"}, {"id": "2", "name": "MENU"}]
 
         client.get_clients.return_value = [
             {"id": "dummy-id", "clientId": "dummy-client"},
@@ -148,7 +149,21 @@ def test_create_resources(self, mock_get_client):
                     client_id="test-id",
                     payload={
                         "name": resource.value,
-                        "scopes": scopes,
+                        "scopes": [{"id": "1", "name": "GET"}],
+                    },
+                    skip_exists=True,
+                )
+            )
+        client.create_client_authz_resource.assert_has_calls(calls)
+
+        calls = []
+        for item in MenuItem:
+            calls.append(
+                call(
+                    client_id="test-id",
+                    payload={
+                        "name": item.value,
+                        "scopes": [{"id": "2", "name": "MENU"}],
                     },
                     skip_exists=True,
                 )

providers/keycloak/tests/unit/keycloak/auth_manager/test_keycloak_auth_manager.py

@@ -34,6 +34,7 @@
     PoolDetails,
     VariableDetails,
 )
+from airflow.api_fastapi.common.types import MenuItem
 from airflow.exceptions import AirflowException
 from airflow.providers.keycloak.auth_manager.constants import (
     CONF_CLIENT_ID_KEY,
@@ -369,5 +370,59 @@ def test_is_authorized_custom_view(
         mock_requests.post.assert_called_once_with(token_url, data=payload, headers=headers)
         assert result == expected
 
+    @pytest.mark.parametrize(
+        "status_code, response, expected",
+        [
+            [
+                200,
+                [{"scopes": ["MENU"], "rsname": "Assets"}, {"scopes": ["MENU"], "rsname": "Connections"}],
+                {MenuItem.ASSETS, MenuItem.CONNECTIONS},
+            ],
+            [200, [{"scopes": ["MENU"], "rsname": "Assets"}], {MenuItem.ASSETS}],
+            [200, [], set()],
+            [403, [{"scopes": ["MENU"], "rsname": "Assets"}], set()],
+        ],
+    )
+    @patch("airflow.providers.keycloak.auth_manager.keycloak_auth_manager.requests")
+    def test_filter_authorized_menu_items(
+        self, mock_requests, status_code, response, expected, auth_manager, user
+    ):
+        mock_requests.post.return_value.status_code = status_code
+        mock_requests.post.return_value.json.return_value = response
+        menu_items = [MenuItem.ASSETS, MenuItem.CONNECTIONS]
+
+        result = auth_manager.filter_authorized_menu_items(menu_items, user=user)
+
+        token_url = auth_manager._get_token_url("server_url", "realm")
+        payload = auth_manager._get_batch_payload(
+            "client_id", [("MENU", MenuItem.ASSETS.value), ("MENU", MenuItem.CONNECTIONS.value)]
+        )
+        headers = auth_manager._get_headers("access_token")
+        mock_requests.post.assert_called_once_with(token_url, data=payload, headers=headers)
+        assert set(result) == expected
+
+    @pytest.mark.parametrize(
+        "status_code",
+        [400, 500],
+    )
+    @patch("airflow.providers.keycloak.auth_manager.keycloak_auth_manager.requests")
+    def test_filter_authorized_menu_items_with_failure(self, mock_requests, status_code, auth_manager, user):
+        resp = Mock()
+        resp.status_code = status_code
+        resp.text = json.dumps({})
+        mock_requests.post.return_value = resp
+
+        menu_items = [MenuItem.ASSETS, MenuItem.CONNECTIONS]
+
+        with pytest.raises(AirflowException):
+            auth_manager.filter_authorized_menu_items(menu_items, user=user)
+
+        token_url = auth_manager._get_token_url("server_url", "realm")
+        payload = auth_manager._get_batch_payload(
+            "client_id", [("MENU", MenuItem.ASSETS.value), ("MENU", MenuItem.CONNECTIONS.value)]
+        )
+        headers = auth_manager._get_headers("access_token")
+        mock_requests.post.assert_called_once_with(token_url, data=payload, headers=headers)
+
     def test_get_cli_commands_return_cli_commands(self, auth_manager):
         assert len(auth_manager.get_cli_commands()) == 1