diff --git a/backend/dataall/modules/datasets/cdk/env_role_dataset_s3_policy.py b/backend/dataall/modules/datasets/cdk/env_role_dataset_s3_policy.py index b59843cd1..8b0b43152 100644 --- a/backend/dataall/modules/datasets/cdk/env_role_dataset_s3_policy.py +++ b/backend/dataall/modules/datasets/cdk/env_role_dataset_s3_policy.py @@ -70,18 +70,16 @@ def _generate_dataset_statements(datasets: List[Dataset]): @staticmethod def _set_allowed_kms_keys_statements(datasets): - allowed_buckets_kms_keys = [] + imported_kms_alias = [] if datasets: + # Datasets belonging to a team and an environment are present in same region and aws account + imported_dataset_resources = [f"arn:aws:kms:{datasets[0].region}:{datasets[0].AwsAccountId}:key/*"] dataset: Dataset for dataset in datasets: if dataset.imported and dataset.importedKmsKey: - key_id = KmsClient(account_id=dataset.AwsAccountId, region=dataset.region).get_key_id( - key_alias=f"alias/{dataset.KmsAlias}" - ) - if key_id: - allowed_buckets_kms_keys.append( - f"arn:aws:kms:{dataset.region}:{dataset.AwsAccountId}:key/{key_id}") - if len(allowed_buckets_kms_keys): + imported_kms_alias.append(f'alias/{dataset.KmsAlias}') + + if len(imported_kms_alias): return iam.PolicyStatement( sid="KMSImportedDatasetAccess", actions=[ @@ -92,6 +90,11 @@ def _set_allowed_kms_keys_statements(datasets): "kms:GenerateDataKey" ], effect=iam.Effect.ALLOW, - resources=allowed_buckets_kms_keys + resources=imported_dataset_resources, + conditions={ + 'ForAnyValue:StringLike': { + 'kms:ResourceAliases' : imported_kms_alias + } + } ) return None