Source code for airflow.providers.amazon.aws.auth_manager.avp.facade

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

import json
from functools import cached_property
from pathlib import Path
from typing import TYPE_CHECKING, Sequence, TypedDict

from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities, get_action_id, get_entity_type
from airflow.providers.amazon.aws.auth_manager.constants import (
    CONF_AVP_POLICY_STORE_ID_KEY,
    CONF_CONN_ID_KEY,
    CONF_REGION_NAME_KEY,
    CONF_SECTION_NAME,
)
from airflow.providers.amazon.aws.hooks.verified_permissions import VerifiedPermissionsHook
from airflow.utils.helpers import prune_dict
from airflow.utils.log.logging_mixin import LoggingMixin

if TYPE_CHECKING:
    from airflow.auth.managers.base_auth_manager import ResourceMethod
    from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser


# Amazon Verified Permissions allows only up to 30 requests per batch_is_authorized call. See
# https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/verifiedpermissions/client/batch_is_authorized.html
[docs]NB_REQUESTS_PER_BATCH = 30
[docs]class IsAuthorizedRequest(TypedDict, total=False): """Represent the parameters of ``is_authorized`` method in AVP facade."""
[docs] method: ResourceMethod
[docs] entity_type: AvpEntities
[docs] entity_id: str | None
[docs] context: dict | None
[docs]class AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin): """ Facade for Amazon Verified Permissions. Used as an intermediate layer between AWS auth manager and Amazon Verified Permissions. """ @cached_property
[docs] def avp_client(self): """Build Amazon Verified Permissions client.""" aws_conn_id = conf.get(CONF_SECTION_NAME, CONF_CONN_ID_KEY) region_name = conf.get(CONF_SECTION_NAME, CONF_REGION_NAME_KEY) return VerifiedPermissionsHook(aws_conn_id=aws_conn_id, region_name=region_name).conn
@cached_property
[docs] def avp_policy_store_id(self): """Get the Amazon Verified Permission policy store ID from config.""" return conf.get_mandatory_value(CONF_SECTION_NAME, CONF_AVP_POLICY_STORE_ID_KEY)
[docs] def is_authorized( self, *, method: ResourceMethod | str, entity_type: AvpEntities, user: AwsAuthManagerUser | None, entity_id: str | None = None, context: dict | None = None, ) -> bool: """ Make an authorization decision against Amazon Verified Permissions. Check whether the user has permissions to access given resource. :param method: the method to perform. The method can also be a string if the action has been defined in a plugin. In that case, the action can be anything (e.g. can_do). See https://github.com/apache/airflow/issues/39144 :param entity_type: the entity type the user accesses :param user: the user :param entity_id: the entity ID the user accesses. If not provided, all entities of the type will be considered. :param context: optional additional context to pass to Amazon Verified Permissions. """ if user is None: return False entity_list = self._get_user_group_entities(user) self.log.debug( "Making authorization request for user=%s, method=%s, entity_type=%s, entity_id=%s", user.get_id(), method, entity_type, entity_id, ) request_params = prune_dict( { "policyStoreId": self.avp_policy_store_id, "principal": {"entityType": get_entity_type(AvpEntities.USER), "entityId": user.get_id()}, "action": { "actionType": get_entity_type(AvpEntities.ACTION), "actionId": get_action_id(entity_type, method), }, "resource": {"entityType": get_entity_type(entity_type), "entityId": entity_id or "*"}, "entities": {"entityList": entity_list}, "context": self._build_context(context), } ) resp = self.avp_client.is_authorized(**request_params) self.log.debug("Authorization response: %s", resp) if len(resp.get("errors", [])) > 0: self.log.error( "Error occurred while making an authorization decision. Errors: %s", resp["errors"] ) raise AirflowException("Error occurred while making an authorization decision.") return resp["decision"] == "ALLOW"
[docs] def get_batch_is_authorized_results( self, *, requests: Sequence[IsAuthorizedRequest], user: AwsAuthManagerUser, ) -> list[dict]: """ Make a batch authorization decision against Amazon Verified Permissions. Return a list of results for each request. :param requests: the list of requests containing the method, the entity_type and the entity ID :param user: the user """ entity_list = self._get_user_group_entities(user) self.log.debug("Making batch authorization request for user=%s, requests=%s", user.get_id(), requests) avp_requests = [self._build_is_authorized_request_payload(request, user) for request in requests] avp_requests_chunks = [ avp_requests[i : i + NB_REQUESTS_PER_BATCH] for i in range(0, len(avp_requests), NB_REQUESTS_PER_BATCH) ] results = [] for avp_requests in avp_requests_chunks: resp = self.avp_client.batch_is_authorized( policyStoreId=self.avp_policy_store_id, requests=avp_requests, entities={"entityList": entity_list}, ) self.log.debug("Authorization response: %s", resp) has_errors = any(len(result.get("errors", [])) > 0 for result in resp["results"]) if has_errors: self.log.error( "Error occurred while making a batch authorization decision. Result: %s", resp["results"] ) raise AirflowException("Error occurred while making a batch authorization decision.") results.extend(resp["results"]) return results
[docs] def batch_is_authorized( self, *, requests: Sequence[IsAuthorizedRequest], user: AwsAuthManagerUser | None, ) -> bool: """ Make a batch authorization decision against Amazon Verified Permissions. Check whether the user has permissions to access all resources. :param requests: the list of requests containing the method, the entity_type and the entity ID :param user: the user """ if user is None: return False results = self.get_batch_is_authorized_results(requests=requests, user=user) return all(result["decision"] == "ALLOW" for result in results)
[docs] def get_batch_is_authorized_single_result( self, *, batch_is_authorized_results: list[dict], request: IsAuthorizedRequest, user: AwsAuthManagerUser, ) -> dict: """ Get a specific authorization result from the output of ``get_batch_is_authorized_results``. :param batch_is_authorized_results: the response from the ``batch_is_authorized`` API :param request: the request information. Used to find the result in the response. :param user: the user """ request_payload = self._build_is_authorized_request_payload(request, user) for result in batch_is_authorized_results: if result["request"] == request_payload: return result self.log.error( "Could not find the authorization result for request %s in results %s.", request_payload, batch_is_authorized_results, ) raise AirflowException("Could not find the authorization result.")
[docs] def is_policy_store_schema_up_to_date(self) -> bool: """Return whether the policy store schema equals the latest version of the schema.""" resp = self.avp_client.get_schema( policyStoreId=self.avp_policy_store_id, ) policy_store_schema = json.loads(resp["schema"]) schema_path = Path(__file__).parents[0] / "schema.json" with open(schema_path) as schema_file: latest_schema = json.load(schema_file) return policy_store_schema == latest_schema
@staticmethod def _get_user_group_entities(user: AwsAuthManagerUser) -> list[dict]: user_entity = { "identifier": {"entityType": get_entity_type(AvpEntities.USER), "entityId": user.get_id()}, "parents": [ {"entityType": get_entity_type(AvpEntities.GROUP), "entityId": group} for group in user.get_groups() ], } group_entities = [ {"identifier": {"entityType": get_entity_type(AvpEntities.GROUP), "entityId": group}} for group in user.get_groups() ] return [user_entity, *group_entities] @staticmethod def _build_context(context: dict | None) -> dict | None: if context is None or len(context) == 0: return None return { "contextMap": context, } def _build_is_authorized_request_payload(self, request: IsAuthorizedRequest, user: AwsAuthManagerUser): """ Build a payload of an individual authorization request that could be sent through the ``batch_is_authorized`` API. :param request: the request information :param user: the user """ return prune_dict( { "principal": {"entityType": get_entity_type(AvpEntities.USER), "entityId": user.get_id()}, "action": { "actionType": get_entity_type(AvpEntities.ACTION), "actionId": get_action_id(request["entity_type"], request["method"]), }, "resource": { "entityType": get_entity_type(request["entity_type"]), "entityId": request.get("entity_id", "*"), }, "context": self._build_context(request.get("context")), } )

Was this entry helpful?