diff --git a/providers/keycloak/docs/auth-manager/manage/permissions.rst b/providers/keycloak/docs/auth-manager/manage/permissions.rst index 7e72cfddf0991..15b0b07407915 100644 --- a/providers/keycloak/docs/auth-manager/manage/permissions.rst +++ b/providers/keycloak/docs/auth-manager/manage/permissions.rst @@ -44,6 +44,7 @@ CLI commands take the following parameters: They also take the following optional parameters: * ``--dry-run``: If set, the command will check the connection to Keycloak and print the actions that would be performed, without actually executing them. +* ``--teams``: Comma-separated list of team names to create team-scoped resources and permissions. Please check the `Keycloak auth manager CLI `_ documentation for more information about accepted parameters. @@ -58,6 +59,12 @@ This command will create scopes, resources and permissions in one-go. airflow keycloak-auth-manager create-all +To create team-scoped resources and permissions (so Keycloak can enforce per-team access) pass ``--teams``: + +.. code-block:: bash + + airflow keycloak-auth-manager create-all --teams team-a,team-b + Step-by-step creation of permissions ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -73,20 +80,45 @@ This command will create resources for certain types of permissions. .. code-block:: bash - airflow keycloak-auth-manager create-resources + airflow keycloak-auth-manager create-resources --teams team-a,team-b Finally, with the command below, we create the permissions using the previously created scopes and resources. .. code-block:: bash - airflow keycloak-auth-manager create-permissions + airflow keycloak-auth-manager create-permissions --teams team-a,team-b This will create -* read-only permissions -* admin permissions -* user permissions -* operations permissions +* read-only permissions (per-team when ``--teams`` is provided) +* admin permissions (global) +* user permissions (per-team when ``--teams`` is provided) +* operations permissions (per-team when ``--teams`` is provided) + +Managing teams with Keycloak +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +When using team-scoped resources, create Keycloak groups that represent teams and attach them to the +team-specific permissions. The CLI provides helpers for this flow: + +.. code-block:: bash + + airflow keycloak-auth-manager create-team team-a + airflow keycloak-auth-manager add-user-to-team user-a team-a + +These commands create a Keycloak group named ``team-a``, set up team-scoped resources and permissions, +and attach team-specific policies to the permissions for that team. +When using team-scoped permissions, the model is: + +* Keycloak group represents the team (````) +* Keycloak roles (Admin/Op/User/Viewer) represent the role within that team +* Team policies require both group membership and role membership +* A separate ``SuperAdmin`` role can be used for global admin access across all teams + +Note: the CLI creates groups, resources, permissions, and policies, but **does not assign roles to users**. +You must assign the appropriate Keycloak roles (Admin/Op/User/Viewer or SuperAdmin) to each user separately. +In multi-team mode, the ``Admin`` role is **team-scoped** (group + role). Only ``SuperAdmin`` grants global +admin access across all teams. More resources about permissions can be found in the official documentation of Keycloak: diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/commands.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/commands.py index 0baaa32d8e793..997b90ae02686 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/commands.py +++ b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/cli/commands.py @@ -22,6 +22,7 @@ from typing import get_args from keycloak import KeycloakAdmin, KeycloakError +from keycloak.exceptions import KeycloakGetError, KeycloakPostError, raise_error_from_response from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod from airflow.api_fastapi.common.types import MenuItem @@ -46,6 +47,30 @@ ) log = logging.getLogger(__name__) +TEAM_SCOPED_RESOURCE_NAMES = { + KeycloakResource.DAG.value, + KeycloakResource.CONNECTION.value, + KeycloakResource.VARIABLE.value, + KeycloakResource.POOL.value, +} +GLOBAL_SCOPED_RESOURCE_NAMES = { + KeycloakResource.ASSET.value, + KeycloakResource.ASSET_ALIAS.value, + KeycloakResource.CONFIGURATION.value, +} +TEAM_MENU_ITEMS = { + MenuItem.DAGS, + MenuItem.ASSETS, + MenuItem.DOCS, +} +TEAM_ADMIN_MENU_ITEMS = TEAM_MENU_ITEMS | { + MenuItem.CONNECTIONS, + MenuItem.POOLS, + MenuItem.VARIABLES, + MenuItem.XCOMS, +} +TEAM_ROLE_NAMES = ("Viewer", "User", "Op", "Admin") +SUPER_ADMIN_ROLE_NAME = "SuperAdmin" def _get_resource_methods() -> list[str]: @@ -90,8 +115,10 @@ def create_resources_command(args): """Create Keycloak auth manager resources in Keycloak.""" client = _get_client(args) client_uuid = _get_client_uuid(args) + teams = _parse_teams(args.teams) + _ensure_multi_team_enabled(teams=teams, command_name="create-resources") - _create_resources(client, client_uuid, _dry_run=args.dry_run) + _create_resources(client, client_uuid, teams=teams, _dry_run=args.dry_run) @cli_utils.action_cli @@ -102,7 +129,14 @@ def create_permissions_command(args): client = _get_client(args) client_uuid = _get_client_uuid(args) - _create_permissions(client, client_uuid, _dry_run=args.dry_run) + teams = _parse_teams(args.teams) + _ensure_multi_team_enabled(teams=teams, command_name="create-permissions") + if teams: + # Role policies are only needed for team-scoped (group+role) authorization. + for role_name in TEAM_ROLE_NAMES: + _ensure_role_policy(client, client_uuid, role_name, _dry_run=args.dry_run) + _ensure_role_policy(client, client_uuid, SUPER_ADMIN_ROLE_NAME, _dry_run=args.dry_run) + _create_permissions(client, client_uuid, teams=teams, _dry_run=args.dry_run) @cli_utils.action_cli @@ -112,10 +146,18 @@ def create_all_command(args): """Create all Keycloak auth manager entities in Keycloak.""" client = _get_client(args) client_uuid = _get_client_uuid(args) + teams = _parse_teams(args.teams) + _ensure_multi_team_enabled(teams=teams, command_name="create-all") _create_scopes(client, client_uuid, _dry_run=args.dry_run) - _create_resources(client, client_uuid, _dry_run=args.dry_run) - _create_permissions(client, client_uuid, _dry_run=args.dry_run) + _create_resources(client, client_uuid, teams=teams, _dry_run=args.dry_run) + _create_group_membership_mapper(client, client_uuid, _dry_run=args.dry_run) + if teams: + # Role policies are only needed for team-scoped (group+role) authorization. + for role_name in TEAM_ROLE_NAMES: + _ensure_role_policy(client, client_uuid, role_name, _dry_run=args.dry_run) + _ensure_role_policy(client, client_uuid, SUPER_ADMIN_ROLE_NAME, _dry_run=args.dry_run) + _create_permissions(client, client_uuid, teams=teams, _dry_run=args.dry_run) def _get_client(args): @@ -145,6 +187,42 @@ def _get_client_uuid(args): return matches[0]["id"] +def _create_group_membership_mapper( + client: KeycloakAdmin, client_uuid: str, *, _dry_run: bool = False +) -> None: + realm = client.connection.realm_name + url = f"admin/realms/{realm}/clients/{client_uuid}/protocol-mappers/models" + data_raw = client.connection.raw_get(url) + mappers = json.loads(data_raw.text) + for mapper in mappers: + if mapper.get("name") == "groups": + return + if mapper.get("protocolMapper") == "oidc-group-membership-mapper": + if mapper.get("config", {}).get("claim.name") == "groups": + return + + if _dry_run: + print("Would create protocol mapper 'groups'.") + return + + payload = { + "name": "groups", + "protocol": "openid-connect", + "protocolMapper": "oidc-group-membership-mapper", + "consentRequired": False, + "config": { + "full.path": "false", + "id.token.claim": "false", + "access.token.claim": "true", + "userinfo.token.claim": "false", + "claim.name": "groups", + "jsonType.label": "String", + }, + } + data_raw = client.connection.raw_post(url, data=json.dumps(payload), max=-1) + raise_error_from_response(data_raw, KeycloakPostError, expected_codes=[201]) + + def _get_scopes_to_create() -> list[dict]: """Get the list of scopes to be created.""" scopes = [{"name": method} for method in _get_resource_methods()] @@ -152,6 +230,19 @@ def _get_scopes_to_create() -> list[dict]: return scopes +def _parse_teams(teams: str | None) -> list[str]: + if not teams: + return [] + return [team.strip() for team in teams.split(",") if team.strip()] + + +def _ensure_multi_team_enabled(*, teams: list[str], command_name: str) -> None: + if not teams: + return + if not conf.getboolean("core", "multi_team", fallback=False): + raise SystemExit(f"{command_name} requires core.multi_team=True when --teams is used.") + + def _preview_scopes(*args, **kwargs): """Preview scopes that would be created.""" scopes = _get_scopes_to_create() @@ -175,6 +266,7 @@ def _create_scopes(client: KeycloakAdmin, client_uuid: str, *, _dry_run: bool = def _get_resources_to_create( client: KeycloakAdmin, client_uuid: str, + teams: list[str], ) -> tuple[list[tuple[str, list[dict]]], list[tuple[str, list[dict]]]]: """ Get the list of resources to be created. @@ -194,14 +286,24 @@ def _get_resources_to_create( ] standard_resources = [(resource.value, scopes) for resource in KeycloakResource] + + if teams: + existing = {resource_name for resource_name, _ in standard_resources} + for team in teams: + for resource_name in TEAM_SCOPED_RESOURCE_NAMES: + name = f"{resource_name}:{team}" + if name in existing: + continue + standard_resources.append((name, scopes)) + existing.add(name) menu_resources = [(item.value, menu_scopes) for item in MenuItem] return standard_resources, menu_resources -def _preview_resources(client: KeycloakAdmin, client_uuid: str): +def _preview_resources(client: KeycloakAdmin, client_uuid: str, teams: list[str]): """Preview resources that would be created.""" - standard_resources, menu_resources = _get_resources_to_create(client, client_uuid) + standard_resources, menu_resources = _get_resources_to_create(client, client_uuid, teams=teams) print("Resources to be created:") if standard_resources: @@ -216,9 +318,9 @@ def _preview_resources(client: KeycloakAdmin, client_uuid: str): @dry_run_preview(_preview_resources) -def _create_resources(client: KeycloakAdmin, client_uuid: str, *, _dry_run: bool = False): +def _create_resources(client: KeycloakAdmin, client_uuid: str, *, teams: list[str], _dry_run: bool = False): """Create Keycloak resources.""" - standard_resources, menu_resources = _get_resources_to_create(client, client_uuid) + standard_resources, menu_resources = _get_resources_to_create(client, client_uuid, teams=teams) for resource_name, scopes in standard_resources: client.create_client_authz_resource( @@ -244,39 +346,118 @@ def _create_resources(client: KeycloakAdmin, client_uuid: str, *, _dry_run: bool print("Resources created successfully.") -def _get_permissions_to_create(client: KeycloakAdmin, client_uuid: str) -> list[dict]: +def _get_permissions_to_create( + client: KeycloakAdmin, + client_uuid: str, + teams: list[str], + *, + include_global_admin: bool = True, +) -> list[dict]: """ Get the actual permissions to be created with filtered scopes/resources. Returns a list of permission descriptors with actual filtered data. """ - perm_configs = [ - { - "name": "ReadOnly", - "type": "scope-based", - "scope_names": ["GET", "MENU", "LIST"], - }, - { - "name": "Admin", - "type": "scope-based", - "scope_names": _get_extended_resource_methods() + ["LIST"], - }, - { - "name": "User", - "type": "resource-based", - "resources": [KeycloakResource.DAG.value, KeycloakResource.ASSET.value], - }, - { - "name": "Op", - "type": "resource-based", - "resources": [ - KeycloakResource.CONNECTION.value, - KeycloakResource.POOL.value, - KeycloakResource.VARIABLE.value, - KeycloakResource.BACKFILL.value, - ], - }, - ] + if not teams: + perm_configs = [ + { + "name": "ReadOnly", + "type": "scope-based", + "scope_names": ["GET", "MENU", "LIST"], + }, + { + "name": "Admin", + "type": "scope-based", + "scope_names": _get_extended_resource_methods() + ["LIST"], + }, + { + "name": "User", + "type": "resource-based", + "resources": [KeycloakResource.DAG.value, KeycloakResource.ASSET.value], + }, + { + "name": "Op", + "type": "resource-based", + "resources": [ + KeycloakResource.CONNECTION.value, + KeycloakResource.POOL.value, + KeycloakResource.VARIABLE.value, + KeycloakResource.BACKFILL.value, + ], + }, + ] + else: + perm_configs = [] + for team in teams: + perm_configs.extend( + [ + { + "name": f"Admin-{team}", + "type": "scope-based", + "scope_names": _get_extended_resource_methods() + ["LIST"], + "resources": [f"{resource}:{team}" for resource in TEAM_SCOPED_RESOURCE_NAMES], + }, + { + "name": f"ReadOnly-{team}", + "type": "scope-based", + "scope_names": ["GET", "LIST"], + "resources": [ + f"{KeycloakResource.DAG.value}:{team}", + ], + }, + { + "name": f"User-{team}", + "type": "resource-based", + "resources": [ + f"{KeycloakResource.DAG.value}:{team}", + ], + }, + { + "name": f"Op-{team}", + "type": "resource-based", + "resources": [ + f"{KeycloakResource.CONNECTION.value}:{team}", + f"{KeycloakResource.POOL.value}:{team}", + f"{KeycloakResource.VARIABLE.value}:{team}", + ], + }, + ] + ) + if include_global_admin: + perm_configs.append( + { + "name": "Admin", + "type": "scope-based", + "scope_names": _get_extended_resource_methods() + ["LIST"], + "resources": [ + f"{resource}:{team}" for team in teams for resource in TEAM_SCOPED_RESOURCE_NAMES + ], + } + ) + perm_configs.append( + { + "name": "GlobalList", + "type": "scope-based", + "scope_names": ["LIST"], + "resources": list(TEAM_SCOPED_RESOURCE_NAMES) + list(GLOBAL_SCOPED_RESOURCE_NAMES), + } + ) + perm_configs.append( + { + "name": "ViewAccess", + "type": "scope-based", + "scope_names": ["GET"], + "resources": [KeycloakResource.VIEW.value], + } + ) + perm_configs.append( + { + "name": "MenuAccess", + "type": "scope-based", + "scope_names": ["MENU"], + "resources": [item.value for item in MenuItem], + } + ) all_scopes = client.get_client_authz_scopes(client_uuid) all_resources = client.get_client_authz_resources(client_uuid) @@ -291,6 +472,13 @@ def _get_permissions_to_create(client: KeycloakAdmin, client_uuid: str) -> list[ filtered_scope_names = [s["name"] for s in all_scopes if s["name"] in config["scope_names"]] perm["scope_ids"] = filtered_scope_ids perm["scope_names"] = filtered_scope_names + if "resources" in config: + filtered_resource_ids = [r["_id"] for r in all_resources if r["name"] in config["resources"]] + filtered_resource_names = [ + r["name"] for r in all_resources if r["name"] in config["resources"] + ] + perm["resource_ids"] = filtered_resource_ids + perm["resource_names"] = filtered_resource_names else: # resource-based # Filter to get actual resource IDs that exist and match filtered_resource_ids = [r["_id"] for r in all_resources if r["name"] in config["resources"]] @@ -302,15 +490,17 @@ def _get_permissions_to_create(client: KeycloakAdmin, client_uuid: str) -> list[ return result -def _preview_permissions(client: KeycloakAdmin, client_uuid: str): +def _preview_permissions(client: KeycloakAdmin, client_uuid: str, teams: list[str]): """Preview permissions that would be created.""" - permissions = _get_permissions_to_create(client, client_uuid) + permissions = _get_permissions_to_create(client, client_uuid, teams=teams) print("Permissions to be created:") for perm in permissions: if perm["type"] == "scope-based": scope_names = ", ".join(perm["scope_names"]) - print(f" - {perm['name']} (type: scope-based, scopes: {scope_names})") + resource_names = ", ".join(perm.get("resource_names", [])) + resource_suffix = f", resources: {resource_names}" if resource_names else "" + print(f" - {perm['name']} (type: scope-based, scopes: {scope_names}{resource_suffix})") else: # resource-based resource_names = ", ".join(perm["resource_names"]) print(f" - {perm['name']} (type: resource-based, resources: {resource_names})") @@ -318,27 +508,47 @@ def _preview_permissions(client: KeycloakAdmin, client_uuid: str): @dry_run_preview(_preview_permissions) -def _create_permissions(client: KeycloakAdmin, client_uuid: str, *, _dry_run: bool = False): +def _create_permissions( + client: KeycloakAdmin, + client_uuid: str, + *, + teams: list[str], + include_global_admin: bool = True, + _dry_run: bool = False, +): """Create Keycloak permissions.""" - permissions = _get_permissions_to_create(client, client_uuid) + permissions = _get_permissions_to_create( + client, client_uuid, teams=teams, include_global_admin=include_global_admin + ) for perm in permissions: if perm["type"] == "scope-based": - _create_scope_based_permission(client, client_uuid, perm["name"], perm["scope_ids"]) + _create_scope_based_permission( + client, client_uuid, perm["name"], perm["scope_ids"], perm.get("resource_ids", []) + ) else: # resource-based _create_resource_based_permission(client, client_uuid, perm["name"], perm["resource_ids"]) print("Permissions created successfully.") -def _create_scope_based_permission(client: KeycloakAdmin, client_uuid: str, name: str, scope_ids: list[str]): +def _create_scope_based_permission( + client: KeycloakAdmin, + client_uuid: str, + name: str, + scope_ids: list[str], + resource_ids: list[str] | None = None, + decision_strategy: str = "UNANIMOUS", +): payload = { "name": name, "type": "scope", "logic": "POSITIVE", - "decisionStrategy": "UNANIMOUS", + "decisionStrategy": decision_strategy, "scopes": scope_ids, } + if resource_ids: + payload["resources"] = resource_ids try: client.create_client_authz_scope_permission( @@ -367,3 +577,625 @@ def _create_resource_based_permission( payload=payload, skip_exists=True, ) + + +def _ensure_scope_permission( + client: KeycloakAdmin, + client_uuid: str, + *, + name: str, + scope_names: list[str], + resource_names: list[str], + decision_strategy: str = "UNANIMOUS", + _dry_run: bool = False, +) -> None: + if _dry_run: + print(f"Would create scope permission '{name}'.") + return + + permissions = client.get_client_authz_permissions(client_uuid) + if any(perm.get("name") == name for perm in permissions): + return + + scopes = client.get_client_authz_scopes(client_uuid) + resources = client.get_client_authz_resources(client_uuid) + scope_ids = [s["id"] for s in scopes if s["name"] in scope_names] + resource_ids = [r["_id"] for r in resources if r["name"] in resource_names] + _create_scope_based_permission( + client, + client_uuid, + name, + scope_ids, + resource_ids, + decision_strategy=decision_strategy, + ) + + +def _update_admin_permission_resources( + client: KeycloakAdmin, client_uuid: str, *, _dry_run: bool = False +) -> None: + if _dry_run: + print("Would update permission 'Admin' with team-scoped and global resources.") + return + + permissions = client.get_client_authz_permissions(client_uuid) + match = next((perm for perm in permissions if perm.get("name") == "Admin"), None) + if not match: + return + + permission_id = match["id"] + scopes = client.get_client_authz_scopes(client_uuid) + resources = client.get_client_authz_resources(client_uuid) + scope_names = _get_extended_resource_methods() + ["LIST"] + scope_ids = [s["id"] for s in scopes if s["name"] in scope_names] + + resource_ids = [ + r["_id"] + for r in resources + if any(r["name"].startswith(f"{resource}:") for resource in TEAM_SCOPED_RESOURCE_NAMES) + or r["name"] in GLOBAL_SCOPED_RESOURCE_NAMES + ] + + policy_ids = _get_permission_policy_ids(client, client_uuid, permission_id) + payload = { + "id": permission_id, + "name": "Admin", + "type": "scope", + "logic": "POSITIVE", + "decisionStrategy": "UNANIMOUS", + "scopes": scope_ids, + "resources": resource_ids, + "policies": policy_ids, + } + client.update_client_authz_scope_permission( + payload=payload, client_id=client_uuid, scope_id=permission_id + ) + + +@cli_utils.action_cli +@providers_configuration_loaded +@dry_run_message_wrap +def create_team_command(args): + """Create team resources, permissions, and Keycloak group.""" + client = _get_client(args) + client_uuid = _get_client_uuid(args) + team = args.team + _ensure_multi_team_enabled(teams=[team], command_name="create-team") + + _create_resources(client, client_uuid, teams=[team], _dry_run=args.dry_run) + _create_group_membership_mapper(client, client_uuid, _dry_run=args.dry_run) + _create_permissions(client, client_uuid, teams=[team], include_global_admin=False, _dry_run=args.dry_run) + _ensure_group(client, team, _dry_run=args.dry_run) + _ensure_team_policies(client, client_uuid, team, _dry_run=args.dry_run) + _attach_team_permissions(client, client_uuid, team, _dry_run=args.dry_run) + _attach_team_menu_permissions(client, client_uuid, team, _dry_run=args.dry_run) + _attach_superadmin_permissions(client, client_uuid, team, _dry_run=args.dry_run) + _update_admin_permission_resources(client, client_uuid, _dry_run=args.dry_run) + + +@cli_utils.action_cli +@providers_configuration_loaded +@dry_run_message_wrap +def add_user_to_team_command(args): + """Add a user to a Keycloak team group.""" + client = _get_client(args) + team = args.team + username = args.target_username + _ensure_multi_team_enabled(teams=[team], command_name="add-user-to-team") + + dry_run = getattr(args, "dry_run", False) + _ensure_group(client, team, _dry_run=dry_run) + _add_user_to_group(client, username=username, team=team, _dry_run=dry_run) + + +def _ensure_team_policies( + client: KeycloakAdmin, client_uuid: str, team: str, *, _dry_run: bool = False +) -> None: + _ensure_group_policy( + client, + client_uuid, + team, + _dry_run=_dry_run, + ) + for role_name in TEAM_ROLE_NAMES: + _ensure_aggregate_policy( + client, + client_uuid, + _team_role_policy_name(team, role_name), + [ + (_team_group_policy_name(team), "group"), + (_role_policy_name(role_name), "role"), + ], + _dry_run=_dry_run, + ) + + +def _attach_team_permissions( + client: KeycloakAdmin, client_uuid: str, team: str, *, _dry_run: bool = False +) -> None: + team_dag_resources = [ + f"{KeycloakResource.DAG.value}:{team}", + ] + team_scoped_resources = [f"{resource}:{team}" for resource in sorted(TEAM_SCOPED_RESOURCE_NAMES)] + + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name=f"ReadOnly-{team}", + policy_name=_team_role_policy_name(team, "Viewer"), + scope_names=["GET", "LIST"], + resource_names=team_dag_resources, + _dry_run=_dry_run, + ) + for role_name in ("User", "Op", "Admin"): + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name=f"ReadOnly-{team}", + policy_name=_team_role_policy_name(team, role_name), + scope_names=["GET", "LIST"], + resource_names=team_dag_resources, + _dry_run=_dry_run, + ) + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name=f"Admin-{team}", + policy_name=_team_role_policy_name(team, "Admin"), + scope_names=_get_extended_resource_methods() + ["LIST"], + resource_names=team_scoped_resources, + _dry_run=_dry_run, + ) + _attach_policy_to_resource_permission( + client, + client_uuid, + permission_name=f"User-{team}", + policy_name=_team_role_policy_name(team, "User"), + resource_names=team_dag_resources, + _dry_run=_dry_run, + ) + _attach_policy_to_resource_permission( + client, + client_uuid, + permission_name=f"Op-{team}", + policy_name=_team_role_policy_name(team, "Op"), + resource_names=[ + f"{KeycloakResource.CONNECTION.value}:{team}", + f"{KeycloakResource.POOL.value}:{team}", + f"{KeycloakResource.VARIABLE.value}:{team}", + ], + _dry_run=_dry_run, + ) + for role_name in TEAM_ROLE_NAMES: + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name="ViewAccess", + policy_name=_team_role_policy_name(team, role_name), + scope_names=["GET"], + resource_names=[KeycloakResource.VIEW.value], + decision_strategy="AFFIRMATIVE", + _dry_run=_dry_run, + ) + + _attach_global_list_permissions(client, client_uuid, _dry_run=_dry_run) + + +def _attach_global_list_permissions( + client: KeycloakAdmin, client_uuid: str, *, _dry_run: bool = False +) -> None: + resource_names = list(TEAM_SCOPED_RESOURCE_NAMES) + list(GLOBAL_SCOPED_RESOURCE_NAMES) + for role_name in (*TEAM_ROLE_NAMES, SUPER_ADMIN_ROLE_NAME): + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name="GlobalList", + policy_name=_role_policy_name(role_name), + scope_names=["LIST"], + resource_names=resource_names, + decision_strategy="AFFIRMATIVE", + _dry_run=_dry_run, + ) + + +def _attach_team_menu_permissions( + client: KeycloakAdmin, client_uuid: str, team: str, *, _dry_run: bool = False +) -> None: + menu_permission_name = f"MenuAccess-{team}" + menu_admin_permission_name = f"MenuAccess-Admin-{team}" + team_menu_resources = [item.value for item in sorted(TEAM_MENU_ITEMS, key=lambda item: item.value)] + team_admin_menu_resources = [ + item.value for item in sorted(TEAM_ADMIN_MENU_ITEMS, key=lambda item: item.value) + ] + _ensure_scope_permission( + client, + client_uuid, + name=menu_permission_name, + scope_names=["MENU"], + resource_names=team_menu_resources, + decision_strategy="AFFIRMATIVE", + _dry_run=_dry_run, + ) + _ensure_scope_permission( + client, + client_uuid, + name=menu_admin_permission_name, + scope_names=["MENU"], + resource_names=team_admin_menu_resources, + decision_strategy="AFFIRMATIVE", + _dry_run=_dry_run, + ) + for role_name in TEAM_ROLE_NAMES: + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name=menu_permission_name, + policy_name=_team_role_policy_name(team, role_name), + scope_names=["MENU"], + resource_names=team_menu_resources, + decision_strategy="AFFIRMATIVE", + _dry_run=_dry_run, + ) + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name=menu_admin_permission_name, + policy_name=_team_role_policy_name(team, "Admin"), + scope_names=["MENU"], + resource_names=team_admin_menu_resources, + decision_strategy="AFFIRMATIVE", + _dry_run=_dry_run, + ) + + +def _attach_superadmin_permissions( + client: KeycloakAdmin, client_uuid: str, team: str, *, _dry_run: bool = False +) -> None: + team_scoped_resources = [f"{resource}:{team}" for resource in sorted(TEAM_SCOPED_RESOURCE_NAMES)] + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name="Admin", + policy_name=_role_policy_name(SUPER_ADMIN_ROLE_NAME), + scope_names=_get_extended_resource_methods() + ["LIST"], + resource_names=team_scoped_resources, + _dry_run=_dry_run, + ) + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name="ViewAccess", + policy_name=_role_policy_name(SUPER_ADMIN_ROLE_NAME), + scope_names=["GET"], + resource_names=[KeycloakResource.VIEW.value], + decision_strategy="AFFIRMATIVE", + _dry_run=_dry_run, + ) + _attach_policy_to_scope_permission( + client, + client_uuid, + permission_name="MenuAccess", + policy_name=_role_policy_name(SUPER_ADMIN_ROLE_NAME), + scope_names=["MENU"], + resource_names=[item.value for item in sorted(MenuItem, key=lambda item: item.value)], + decision_strategy="AFFIRMATIVE", + _dry_run=_dry_run, + ) + + +def _team_group_name(team: str) -> str: + return team + + +def _team_group_policy_name(team: str) -> str: + return f"Allow-Team-{team}" + + +def _role_policy_name(role_name: str) -> str: + return f"Allow-{role_name}" + + +def _team_role_policy_name(team: str, role_name: str) -> str: + return f"Allow-{role_name}-{team}" + + +def _ensure_group(client: KeycloakAdmin, team: str, *, _dry_run: bool = False) -> dict | None: + group_name = _team_group_name(team) + group_path = f"/{group_name}" + try: + group = client.get_group_by_path(group_path) + if group: + return group + except KeycloakError: + pass + + if _dry_run: + print(f"Would create group '{group_name}'.") + return None + + group_id = client.create_group(payload={"name": group_name}, skip_exists=True) + if not group_id: + group = client.get_group_by_path(group_path) + return group + return {"id": group_id, "name": group_name, "path": group_path} + + +def _ensure_group_policy( + client: KeycloakAdmin, client_uuid: str, team: str, *, _dry_run: bool = False +) -> None: + group_name = _team_group_name(team) + group_path = f"/{group_name}" + policy_name = _team_group_policy_name(team) + try: + group = client.get_group_by_path(group_path) + except KeycloakError: + group = None + if not group: + raise ValueError(f"Group '{group_name}' not found.") + + if _get_policy_id(client, client_uuid, policy_name, policy_type="group"): + return + + payload = { + "name": policy_name, + "type": "group", + "logic": "POSITIVE", + "decisionStrategy": "UNANIMOUS", + "groups": [{"id": group["id"], "path": group_path}], + } + + if _dry_run: + print(f"Would create group policy '{policy_name}'.") + return + + url = _policy_url(client, client_uuid, policy_type="group") + data_raw = client.connection.raw_post(url, data=json.dumps(payload), max=-1, permission=False) + try: + raise_error_from_response(data_raw, KeycloakPostError, expected_codes=[201]) + except KeycloakPostError as exc: + if exc.response_body: + error = json.loads(exc.response_body.decode("utf-8")) + if "Conflicting policy" in error.get("error_description", ""): + return + raise + + +def _policy_url(client: KeycloakAdmin, client_uuid: str, *, policy_type: str | None = None) -> str: + realm = client.connection.realm_name + if policy_type: + return f"admin/realms/{realm}/clients/{client_uuid}/authz/resource-server/policy/{policy_type}" + return f"admin/realms/{realm}/clients/{client_uuid}/authz/resource-server/policy" + + +def _get_policy_id( + client: KeycloakAdmin, client_uuid: str, policy_name: str, *, policy_type: str | None = None +) -> str | None: + url = _policy_url(client, client_uuid, policy_type=policy_type) + data_raw = client.connection.raw_get(url) + policies = json.loads(data_raw.text) + match = next((policy for policy in policies if policy.get("name") == policy_name), None) + return match.get("id") if match else None + + +def _get_role_id(client: KeycloakAdmin, client_uuid: str, role_name: str) -> str: + try: + role = client.get_realm_role(role_name) + return role["id"] + except KeycloakGetError: + role = client.get_client_role(client_id=client_uuid, role_name=role_name) + return role["id"] + + +def _ensure_role_policy( + client: KeycloakAdmin, client_uuid: str, role_name: str, *, _dry_run: bool = False +) -> None: + policy_name = _role_policy_name(role_name) + if _get_policy_id(client, client_uuid, policy_name, policy_type="role"): + return + + role_id = _get_role_id(client, client_uuid, role_name) + payload = { + "name": policy_name, + "type": "role", + "logic": "POSITIVE", + "decisionStrategy": "UNANIMOUS", + "roles": [{"id": role_id}], + } + + if _dry_run: + print(f"Would create role policy '{policy_name}'.") + return + + try: + client.create_client_authz_role_based_policy(client_id=client_uuid, payload=payload, skip_exists=True) + except KeycloakError as exc: + if exc.response_body: + error = json.loads(exc.response_body.decode("utf-8")) + if "Conflicting policy" in error.get("error_description", ""): + return + raise + + +def _ensure_aggregate_policy( + client: KeycloakAdmin, + client_uuid: str, + policy_name: str, + policy_refs: list[tuple[str, str | None]], + *, + _dry_run: bool = False, +) -> None: + if _get_policy_id(client, client_uuid, policy_name, policy_type="aggregate"): + return + + policy_ids = [] + for name, policy_type in policy_refs: + # Aggregate policy enforces group+role; missing inputs should fail fast. + policy_id = _get_policy_id(client, client_uuid, name, policy_type=policy_type) + if not policy_id: + policy_label = f"{policy_type} policy '{name}'" if policy_type else f"policy '{name}'" + raise ValueError(f"{policy_label} not found.") + policy_ids.append(policy_id) + + payload = { + "name": policy_name, + "type": "aggregate", + "logic": "POSITIVE", + "decisionStrategy": "UNANIMOUS", + "policies": policy_ids, + } + + if _dry_run: + print(f"Would create aggregate policy '{policy_name}'.") + return + + url = _policy_url(client, client_uuid, policy_type="aggregate") + data_raw = client.connection.raw_post(url, data=json.dumps(payload), max=-1, permission=False) + try: + raise_error_from_response(data_raw, KeycloakPostError, expected_codes=[201]) + except KeycloakPostError as exc: + if exc.response_body: + error = json.loads(exc.response_body.decode("utf-8")) + if "Conflicting policy" in error.get("error_description", ""): + return + raise + + +def _attach_policy_to_scope_permission( + client: KeycloakAdmin, + client_uuid: str, + *, + permission_name: str, + policy_name: str, + scope_names: list[str], + resource_names: list[str], + decision_strategy: str = "UNANIMOUS", + _dry_run: bool = False, +) -> None: + if _dry_run: + print(f"Would attach policy '{policy_name}' to permission '{permission_name}'.") + return + + permissions = client.get_client_authz_permissions(client_uuid) + match = next((perm for perm in permissions if perm.get("name") == permission_name), None) + if not match: + raise ValueError(f"Permission '{permission_name}' not found.") + + permission_id = match["id"] + policy_id = _get_policy_id(client, client_uuid, policy_name) + if not policy_id: + raise ValueError(f"Policy '{policy_name}' not found.") + + existing_policy_ids = _get_permission_policy_ids(client, client_uuid, permission_id) + policy_ids = list(dict.fromkeys([*existing_policy_ids, policy_id])) + + scopes = client.get_client_authz_scopes(client_uuid) + resources = client.get_client_authz_resources(client_uuid) + scope_ids = [s["id"] for s in scopes if s["name"] in scope_names] + resource_ids = [r["_id"] for r in resources if r["name"] in resource_names] + + payload = { + "id": permission_id, + "name": permission_name, + "type": "scope", + "logic": "POSITIVE", + "decisionStrategy": decision_strategy, + "scopes": scope_ids, + "resources": resource_ids, + "policies": policy_ids, + } + client.update_client_authz_scope_permission( + payload=payload, client_id=client_uuid, scope_id=permission_id + ) + + +def _get_permission_policy_ids(client: KeycloakAdmin, client_uuid: str, permission_id: str) -> list[str]: + realm = client.connection.realm_name + url = ( + f"admin/realms/{realm}/clients/{client_uuid}/authz/resource-server/permission/scope/" + f"{permission_id}/associatedPolicies" + ) + data_raw = client.connection.raw_get(url) + policies = json.loads(data_raw.text) + return [policy.get("id") for policy in policies if policy.get("id")] + + +def _attach_policy_to_resource_permission( + client: KeycloakAdmin, + client_uuid: str, + *, + permission_name: str, + policy_name: str, + resource_names: list[str], + decision_strategy: str = "UNANIMOUS", + _dry_run: bool = False, +) -> None: + if _dry_run: + print(f"Would attach policy '{policy_name}' to permission '{permission_name}'.") + return + + permissions = client.get_client_authz_permissions(client_uuid) + match = next((perm for perm in permissions if perm.get("name") == permission_name), None) + if not match: + raise ValueError(f"Permission '{permission_name}' not found.") + + permission_id = match["id"] + policy_id = _get_policy_id(client, client_uuid, policy_name) + if not policy_id: + raise ValueError(f"Policy '{policy_name}' not found.") + + existing_policy_ids = _get_resource_permission_policy_ids(client, client_uuid, permission_id) + policy_ids = list(dict.fromkeys([*existing_policy_ids, policy_id])) + + resources = client.get_client_authz_resources(client_uuid) + resource_ids = [r["_id"] for r in resources if r["name"] in resource_names] + + payload = { + "id": permission_id, + "name": permission_name, + "type": "resource", + "logic": "POSITIVE", + "decisionStrategy": decision_strategy, + "resources": resource_ids, + "scopes": [], + "policies": policy_ids, + } + client.update_client_authz_resource_permission( + payload=payload, client_id=client_uuid, resource_id=permission_id + ) + + +def _get_resource_permission_policy_ids( + client: KeycloakAdmin, client_uuid: str, permission_id: str +) -> list[str]: + realm = client.connection.realm_name + url = ( + f"admin/realms/{realm}/clients/{client_uuid}/authz/resource-server/permission/resource/" + f"{permission_id}/associatedPolicies" + ) + data_raw = client.connection.raw_get(url) + policies = json.loads(data_raw.text) + return [policy.get("id") for policy in policies if policy.get("id")] + + +def _add_user_to_group(client: KeycloakAdmin, *, username: str, team: str, _dry_run: bool = False) -> None: + group_name = _team_group_name(team) + group_path = f"/{group_name}" + group = client.get_group_by_path(group_path) + if not group: + raise ValueError(f"Group '{group_name}' not found.") + + users = client.get_users(query={"username": username}) + user = next((u for u in users if u.get("username") == username), None) + if not user: + raise ValueError(f"User '{username}' not found.") + + existing_groups = client.get_user_groups(user_id=user["id"]) + if any(g.get("id") == group["id"] for g in existing_groups): + print(f"User '{username}' is already in group '{group_name}'.") + return + + if _dry_run: + print(f"Would add user '{username}' to group '{group_name}'.") + return + + client.group_user_add(user_id=user["id"], group_id=group["id"]) diff --git a/providers/keycloak/src/airflow/providers/keycloak/cli/definition.py b/providers/keycloak/src/airflow/providers/keycloak/cli/definition.py index 218f272c07915..dbeb71ca9086d 100644 --- a/providers/keycloak/src/airflow/providers/keycloak/cli/definition.py +++ b/providers/keycloak/src/airflow/providers/keycloak/cli/definition.py @@ -57,6 +57,9 @@ def __call__(self, parser, namespace, values, option_string=None): ("--user-realm",), help="Realm name where the user used to create resources is", default="master" ) ARG_CLIENT_ID = Arg(("--client-id",), help="ID of the client used to create resources", default="admin-cli") +ARG_TEAMS = Arg(("--teams",), help="Comma-separated list of team names") +ARG_TEAM = Arg(("team",), help="Team name") +ARG_TARGET_USER = Arg(("target_username",), help="Username to add to the team") ARG_DRY_RUN = Arg( ("--dry-run",), help="Perform a dry run without creating any resources", @@ -81,7 +84,7 @@ def __call__(self, parser, namespace, values, option_string=None): func=lazy_load_command( "airflow.providers.keycloak.auth_manager.cli.commands.create_resources_command" ), - args=(ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID, ARG_DRY_RUN), + args=(ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID, ARG_TEAMS, ARG_DRY_RUN), ), ActionCommand( name="create-permissions", @@ -89,13 +92,27 @@ def __call__(self, parser, namespace, values, option_string=None): func=lazy_load_command( "airflow.providers.keycloak.auth_manager.cli.commands.create_permissions_command" ), - args=(ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID, ARG_DRY_RUN), + args=(ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID, ARG_TEAMS, ARG_DRY_RUN), ), ActionCommand( name="create-all", help="Create all entities (scopes, resources and permissions) in Keycloak", func=lazy_load_command("airflow.providers.keycloak.auth_manager.cli.commands.create_all_command"), - args=(ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID, ARG_DRY_RUN), + args=(ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID, ARG_TEAMS, ARG_DRY_RUN), + ), + ActionCommand( + name="create-team", + help="Create Keycloak team group, resources, and permissions", + func=lazy_load_command("airflow.providers.keycloak.auth_manager.cli.commands.create_team_command"), + args=(ARG_TEAM, ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID, ARG_DRY_RUN), + ), + ActionCommand( + name="add-user-to-team", + help="Add a Keycloak user to a team group", + func=lazy_load_command( + "airflow.providers.keycloak.auth_manager.cli.commands.add_user_to_team_command" + ), + args=(ARG_TARGET_USER, ARG_TEAM, ARG_USERNAME, ARG_PASSWORD, ARG_USER_REALM, ARG_CLIENT_ID), ), ) diff --git a/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_commands.py b/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_commands.py index c8fa22d81f52d..d80649c32ca70 100644 --- a/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_commands.py +++ b/providers/keycloak/tests/unit/keycloak/auth_manager/cli/test_commands.py @@ -24,11 +24,15 @@ from airflow.api_fastapi.common.types import MenuItem from airflow.cli import cli_parser from airflow.providers.keycloak.auth_manager.cli.commands import ( + TEAM_SCOPED_RESOURCE_NAMES, + _get_extended_resource_methods, _get_resource_methods, + add_user_to_team_command, create_all_command, create_permissions_command, create_resources_command, create_scopes_command, + create_team_command, ) from airflow.providers.keycloak.auth_manager.resources import KeycloakResource @@ -136,6 +140,7 @@ def test_create_resources(self, mock_get_client): with conf_vars( { ("keycloak_auth_manager", "client_id"): "test_client_id", + ("core", "multi_team"): "True", } ): create_resources_command(self.arg_parser.parse_args(params)) @@ -170,6 +175,44 @@ def test_create_resources(self, mock_get_client): ) client.create_client_authz_resource.assert_has_calls(calls) + @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client") + def test_create_resources_with_teams(self, mock_get_client): + client = Mock() + mock_get_client.return_value = client + scopes = [{"id": "1", "name": "GET"}, {"id": "2", "name": "MENU"}] + + client.get_clients.return_value = [ + {"id": "dummy-id", "clientId": "dummy-client"}, + {"id": "test-id", "clientId": "test_client_id"}, + ] + client.get_client_authz_scopes.return_value = scopes + + params = [ + "keycloak-auth-manager", + "create-resources", + "--username", + "test", + "--password", + "test", + "--teams", + "team-a", + ] + with conf_vars( + { + ("keycloak_auth_manager", "client_id"): "test_client_id", + ("core", "multi_team"): "True", + } + ): + create_resources_command(self.arg_parser.parse_args(params)) + + expected_team_resources = {f"{name}:team-a" for name in TEAM_SCOPED_RESOURCE_NAMES} + created_resource_names = { + call.kwargs["payload"]["name"] + for call in client.create_client_authz_resource.mock_calls + if "payload" in call.kwargs + } + assert expected_team_resources.issubset(created_resource_names) + @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client") def test_create_permissions( self, @@ -190,6 +233,10 @@ def test_create_permissions( ] client.get_client_authz_scopes.return_value = scopes client.get_client_authz_resources.return_value = resources + client.connection = Mock() + client.connection.raw_get = Mock(return_value=Mock(text="[]")) + client.connection.realm_name = "test-realm" + client.get_realm_role.return_value = {"id": "role-id"} params = [ "keycloak-auth-manager", @@ -202,6 +249,7 @@ def test_create_permissions( with conf_vars( { ("keycloak_auth_manager", "client_id"): "test_client_id", + ("core", "multi_team"): "True", } ): create_permissions_command(self.arg_parser.parse_args(params)) @@ -262,6 +310,322 @@ def test_create_permissions( ] client.create_client_authz_resource_based_permission.assert_has_calls(resource_calls, any_order=True) + @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client") + def test_create_permissions_with_teams(self, mock_get_client): + client = Mock() + mock_get_client.return_value = client + scopes = [ + {"id": "1", "name": "GET"}, + {"id": "2", "name": "MENU"}, + {"id": "3", "name": "LIST"}, + ] + resources = [ + {"_id": "r1", "name": "Dag:team-a"}, + {"_id": "r2", "name": "Connection:team-a"}, + {"_id": "r3", "name": "Pool:team-a"}, + {"_id": "r4", "name": "Variable:team-a"}, + {"_id": "r5", "name": "View"}, + {"_id": "r6", "name": "Dag"}, + {"_id": "r7", "name": "Connection"}, + {"_id": "r8", "name": "Pool"}, + {"_id": "r9", "name": "Variable"}, + {"_id": "r10", "name": "Asset"}, + {"_id": "r11", "name": "AssetAlias"}, + {"_id": "r12", "name": "Configuration"}, + ] + + client.get_clients.return_value = [ + {"id": "dummy-id", "clientId": "dummy-client"}, + {"id": "test-id", "clientId": "test_client_id"}, + ] + client.get_client_authz_scopes.return_value = scopes + client.get_client_authz_resources.return_value = resources + client.connection = Mock() + client.connection.raw_get = Mock(return_value=Mock(text="[]")) + client.connection.realm_name = "test-realm" + client.get_realm_role.return_value = {"id": "role-id"} + + params = [ + "keycloak-auth-manager", + "create-permissions", + "--username", + "test", + "--password", + "test", + "--teams", + "team-a", + ] + with conf_vars( + { + ("keycloak_auth_manager", "client_id"): "test_client_id", + ("core", "multi_team"): "True", + } + ): + create_permissions_command(self.arg_parser.parse_args(params)) + + client.create_client_authz_scope_permission.assert_any_call( + client_id="test-id", + payload={ + "name": "Admin-team-a", + "type": "scope", + "logic": "POSITIVE", + "decisionStrategy": "UNANIMOUS", + "scopes": ["1", "2", "3"], + "resources": ["r1", "r2", "r3", "r4"], + }, + ) + client.create_client_authz_scope_permission.assert_any_call( + client_id="test-id", + payload={ + "name": "GlobalList", + "type": "scope", + "logic": "POSITIVE", + "decisionStrategy": "UNANIMOUS", + "scopes": ["3"], + "resources": ["r6", "r7", "r8", "r9", "r10", "r11", "r12"], + }, + ) + client.create_client_authz_scope_permission.assert_any_call( + client_id="test-id", + payload={ + "name": "ViewAccess", + "type": "scope", + "logic": "POSITIVE", + "decisionStrategy": "UNANIMOUS", + "scopes": ["1"], + "resources": ["r5"], + }, + ) + client.create_client_authz_scope_permission.assert_any_call( + client_id="test-id", + payload={ + "name": "ReadOnly-team-a", + "type": "scope", + "logic": "POSITIVE", + "decisionStrategy": "UNANIMOUS", + "scopes": ["1", "3"], + "resources": ["r1"], + }, + ) + client.create_client_authz_resource_based_permission.assert_any_call( + client_id="test-id", + payload={ + "name": "User-team-a", + "type": "scope", + "logic": "POSITIVE", + "decisionStrategy": "UNANIMOUS", + "resources": ["r1"], + }, + skip_exists=True, + ) + + @patch("airflow.providers.keycloak.auth_manager.cli.commands._update_admin_permission_resources") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._ensure_scope_permission") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._attach_policy_to_resource_permission") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._attach_policy_to_scope_permission") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._ensure_aggregate_policy") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._ensure_group_policy") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._create_permissions") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._create_resources") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._ensure_group") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client") + def test_create_team_command( + self, + mock_get_client, + mock_ensure_group, + mock_create_resources, + mock_create_permissions, + mock_ensure_group_policy, + mock_ensure_aggregate_policy, + mock_attach_policy, + mock_attach_resource_policy, + mock_ensure_scope_permission, + mock_update_admin_permission_resources, + ): + client = Mock() + mock_get_client.return_value = client + client.get_clients.return_value = [ + {"id": "test-id", "clientId": "test_client_id"}, + ] + client.connection = Mock() + client.connection.raw_get = Mock(return_value=Mock(text="[]")) + client.connection.raw_post = Mock( + return_value=Mock(status_code=201, json=Mock(return_value={"message": ""}), text="{}") + ) + client.connection.realm_name = "test-realm" + + params = [ + "keycloak-auth-manager", + "create-team", + "team-a", + "--username", + "test", + "--password", + "test", + ] + with conf_vars( + { + ("keycloak_auth_manager", "client_id"): "test_client_id", + ("core", "multi_team"): "True", + } + ): + create_team_command(self.arg_parser.parse_args(params)) + + mock_ensure_group.assert_called_once_with(client, "team-a", _dry_run=False) + mock_create_resources.assert_called_once_with(client, "test-id", teams=["team-a"], _dry_run=False) + mock_create_permissions.assert_called_once_with( + client, "test-id", teams=["team-a"], include_global_admin=False, _dry_run=False + ) + mock_update_admin_permission_resources.assert_called_once_with(client, "test-id", _dry_run=False) + mock_ensure_group_policy.assert_called_once_with(client, "test-id", "team-a", _dry_run=False) + assert mock_ensure_aggregate_policy.call_count == 4 + mock_attach_policy.assert_any_call( + client, + "test-id", + permission_name="ReadOnly-team-a", + policy_name="Allow-Viewer-team-a", + scope_names=["GET", "LIST"], + resource_names=["Dag:team-a"], + _dry_run=False, + ) + mock_attach_policy.assert_any_call( + client, + "test-id", + permission_name="Admin-team-a", + policy_name="Allow-Admin-team-a", + scope_names=_get_extended_resource_methods() + ["LIST"], + resource_names=[ + "Connection:team-a", + "Dag:team-a", + "Pool:team-a", + "Variable:team-a", + ], + _dry_run=False, + ) + mock_attach_resource_policy.assert_any_call( + client, + "test-id", + permission_name="User-team-a", + policy_name="Allow-User-team-a", + resource_names=["Dag:team-a"], + _dry_run=False, + ) + mock_attach_resource_policy.assert_any_call( + client, + "test-id", + permission_name="Op-team-a", + policy_name="Allow-Op-team-a", + resource_names=[ + "Connection:team-a", + "Pool:team-a", + "Variable:team-a", + ], + _dry_run=False, + ) + mock_attach_policy.assert_any_call( + client, + "test-id", + permission_name="MenuAccess-team-a", + policy_name="Allow-Viewer-team-a", + scope_names=["MENU"], + resource_names=["Assets", "Dags", "Docs"], + decision_strategy="AFFIRMATIVE", + _dry_run=False, + ) + mock_attach_policy.assert_any_call( + client, + "test-id", + permission_name="MenuAccess-Admin-team-a", + policy_name="Allow-Admin-team-a", + scope_names=["MENU"], + resource_names=["Assets", "Connections", "Dags", "Docs", "Pools", "Variables", "XComs"], + decision_strategy="AFFIRMATIVE", + _dry_run=False, + ) + mock_attach_policy.assert_any_call( + client, + "test-id", + permission_name="Admin", + policy_name="Allow-SuperAdmin", + scope_names=_get_extended_resource_methods() + ["LIST"], + resource_names=[ + "Connection:team-a", + "Dag:team-a", + "Pool:team-a", + "Variable:team-a", + ], + _dry_run=False, + ) + mock_attach_policy.assert_any_call( + client, + "test-id", + permission_name="MenuAccess", + policy_name="Allow-SuperAdmin", + scope_names=["MENU"], + resource_names=[item.value for item in sorted(MenuItem, key=lambda item: item.value)], + decision_strategy="AFFIRMATIVE", + _dry_run=False, + ) + mock_attach_policy.assert_any_call( + client, + "test-id", + permission_name="ViewAccess", + policy_name="Allow-Viewer-team-a", + scope_names=["GET"], + resource_names=["View"], + decision_strategy="AFFIRMATIVE", + _dry_run=False, + ) + mock_ensure_scope_permission.assert_any_call( + client, + "test-id", + name="MenuAccess-team-a", + scope_names=["MENU"], + resource_names=["Assets", "Dags", "Docs"], + decision_strategy="AFFIRMATIVE", + _dry_run=False, + ) + mock_ensure_scope_permission.assert_any_call( + client, + "test-id", + name="MenuAccess-Admin-team-a", + scope_names=["MENU"], + resource_names=["Assets", "Connections", "Dags", "Docs", "Pools", "Variables", "XComs"], + decision_strategy="AFFIRMATIVE", + _dry_run=False, + ) + + @patch("airflow.providers.keycloak.auth_manager.cli.commands._add_user_to_group") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._ensure_group") + @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client") + def test_add_user_to_team_command(self, mock_get_client, mock_ensure_group, mock_add_user): + client = Mock() + mock_get_client.return_value = client + client.get_clients.return_value = [ + {"id": "test-id", "clientId": "test_client_id"}, + ] + + params = [ + "keycloak-auth-manager", + "add-user-to-team", + "user-a", + "team-a", + "--username", + "admin", + "--password", + "admin", + ] + with conf_vars( + { + ("keycloak_auth_manager", "client_id"): "test_client_id", + ("core", "multi_team"): "True", + } + ): + add_user_to_team_command(self.arg_parser.parse_args(params)) + + mock_ensure_group.assert_called_once_with(client, "team-a", _dry_run=False) + mock_add_user.assert_called_once_with(client, username="user-a", team="team-a", _dry_run=False) + @patch("airflow.providers.keycloak.auth_manager.cli.commands._create_permissions") @patch("airflow.providers.keycloak.auth_manager.cli.commands._create_resources") @patch("airflow.providers.keycloak.auth_manager.cli.commands._create_scopes") @@ -282,6 +646,12 @@ def test_create_all( {"id": "test-id", "clientId": "test_client_id"}, ] client.get_client_authz_scopes.return_value = scopes + client.connection = Mock() + client.connection.raw_get = Mock(return_value=Mock(text="[]")) + client.connection.raw_post = Mock( + return_value=Mock(status_code=201, json=Mock(return_value={"message": ""}), text="{}") + ) + client.connection.realm_name = "test-realm" params = [ "keycloak-auth-manager", @@ -300,8 +670,8 @@ def test_create_all( client.get_clients.assert_called_once_with() mock_create_scopes.assert_called_once_with(client, "test-id", _dry_run=False) - mock_create_resources.assert_called_once_with(client, "test-id", _dry_run=False) - mock_create_permissions.assert_called_once_with(client, "test-id", _dry_run=False) + mock_create_resources.assert_called_once_with(client, "test-id", teams=[], _dry_run=False) + mock_create_permissions.assert_called_once_with(client, "test-id", teams=[], _dry_run=False) @patch("airflow.providers.keycloak.auth_manager.cli.commands._get_client") def test_create_scopes_dry_run(self, mock_get_client): @@ -424,6 +794,12 @@ def test_create_all_dry_run( {"id": "dummy-id", "clientId": "dummy-client"}, {"id": "test-id", "clientId": "test_client_id"}, ] + client.connection = Mock() + client.connection.raw_get = Mock(return_value=Mock(text="[]")) + client.connection.raw_post = Mock( + return_value=Mock(status_code=201, json=Mock(return_value={"message": ""}), text="{}") + ) + client.connection.realm_name = "test-realm" params = [ "keycloak-auth-manager", @@ -444,5 +820,5 @@ def test_create_all_dry_run( client.get_clients.assert_called_once_with() # In dry-run mode, all helper functions should be called with dry_run=True mock_create_scopes.assert_called_once_with(client, "test-id", _dry_run=True) - mock_create_resources.assert_called_once_with(client, "test-id", _dry_run=True) - mock_create_permissions.assert_called_once_with(client, "test-id", _dry_run=True) + mock_create_resources.assert_called_once_with(client, "test-id", teams=[], _dry_run=True) + mock_create_permissions.assert_called_once_with(client, "test-id", teams=[], _dry_run=True) diff --git a/providers/keycloak/tests/unit/keycloak/cli/test_definition.py b/providers/keycloak/tests/unit/keycloak/cli/test_definition.py index d10a727f69460..fd85e99479aec 100644 --- a/providers/keycloak/tests/unit/keycloak/cli/test_definition.py +++ b/providers/keycloak/tests/unit/keycloak/cli/test_definition.py @@ -48,7 +48,7 @@ def setup_parser(self): self.arg_parser = cli_parser.get_parser() def test_keycloak_auth_manager_cli_commands(self): - assert len(KEYCLOAK_AUTH_MANAGER_COMMANDS) == 4 + assert len(KEYCLOAK_AUTH_MANAGER_COMMANDS) == 6 @pytest.mark.parametrize( "command",