From a273e52f36a61e310a67946a37c2c89f43d5b1a3 Mon Sep 17 00:00:00 2001
From: Tejas Rajopadhye <71188245+TejasRGitHub@users.noreply.github.com>
Date: Tue, 9 Jan 2024 06:25:05 -0600
Subject: [PATCH] Fix when migrating from Manually Created Pivot Role to Auto
 Create Pivot Role (#948)

### Feature or Bugfix
- Bugfix

### Detail

- Changes in the way KMS keys of datasets( S3 buckets ) are specified in
the environment role .

### Testing

- Created an environment with manual pivot role and migrated to auto
create pivot role.
- Checked if the datasets previously present in the environment get
updated properly.
- Imported another dataset on this environment successfully.

### Relates
- https://ouryahoo.atlassian.net/browse/DATA-416

### Security
Please answer the questions below briefly where applicable, or write
`N/A`. Based on
[OWASP 10](https://owasp.org/Top10/en/).

- Does this PR introduce or modify any input fields or queries - this
includes
fetching data from storage outside the application (e.g. a database, an
S3 bucket)? N/A
  - Is the input sanitized?
- What precautions are you taking before deserializing the data you
consume?
  - Is injection prevented by parametrizing queries?
  - Have you ensured no `eval` or similar functions are used?
- Does this PR introduce any functionality or component that requires
authorization? N/A
- How have you ensured it respects the existing AuthN/AuthZ mechanisms?
  - Are you logging failed auth attempts?
- Are you using or adding any cryptographic features? N/A
  - Do you use a standard proven implementations?
  - Are the used keys controlled by the customer? Where are they stored?
- Are you introducing any new policies/roles/users? Yes
- Have you used the least-privilege principle? How? **Yes, only access
to KMS keys which are present on the environment**


By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license.

---------

Co-authored-by: trajopadhye <tejas.rajopadhye@yahooinc.com>
---
 .../cdk/env_role_dataset_s3_policy.py         | 21 +++++++++++--------
 1 file changed, 12 insertions(+), 9 deletions(-)

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