Kms explosion fix

backend/dataall/modules/dataset_sharing/services/data_sharing_service.py

@@ -204,8 +204,7 @@ def revoke_share(cls, engine: Engine, share_uri: str):
                     source_environment=source_environment,
                     target_environment=target_environment,
                     source_env_group=source_env_group,
-                    env_group=env_group,
-                    existing_shared_buckets=existing_shared_buckets
+                    env_group=env_group
                 )
                 log.info(f"Clean up S3 successful = {clean_up_folders}")
 
@@ -219,8 +218,7 @@ def revoke_share(cls, engine: Engine, share_uri: str):
                 source_environment,
                 target_environment,
                 source_env_group,
-                env_group,
-                existing_shared_folders
+                env_group
             )
             log.info(f'revoking s3 buckets succeeded = {revoked_s3_buckets_succeed}')
 

backend/dataall/modules/dataset_sharing/services/share_managers/s3_access_point_share_manager.py

@@ -2,6 +2,7 @@
 import logging
 import json
 import time
+from itertools import count
 
 from dataall.core.environment.db.environment_models import Environment, EnvironmentGroup
 from dataall.base.db import utils
@@ -21,6 +22,7 @@
 ACCESS_POINT_CREATION_RETRIES = 5
 IAM_ACCESS_POINT_ROLE_POLICY = "targetDatasetAccessControlPolicy"
 DATAALL_ALLOW_OWNER_SID = "AllowAllToAdmin"
+DATAALL_ACCESS_POINT_KMS_DECRYPT_SID = "DataAll-Access-Point-KMS-Decrypt"
 
 
 class S3AccessPointShareManager:
@@ -250,7 +252,8 @@ def manage_access_point_and_policy(self):
             access_point_arn = s3_client.create_bucket_access_point(self.bucket_name, self.access_point_name)
             # Access point creation is slow
             retries = 1
-            while not s3_client.get_bucket_access_point_arn(self.access_point_name) and retries < ACCESS_POINT_CREATION_RETRIES:
+            while (not s3_client.get_bucket_access_point_arn(self.access_point_name)
+                   and retries < ACCESS_POINT_CREATION_RETRIES):
                 logger.info(
                     'Waiting 30s for access point creation to complete..'
                 )
@@ -300,7 +303,10 @@ def manage_access_point_and_policy(self):
             )
             exceptions_roleId = [f'{item}:*' for item in SessionHelper.get_role_ids(
                 self.source_account_id,
-                [self.dataset_admin, self.source_env_admin, SessionHelper.get_delegation_role_arn(self.source_account_id)]
+                [
+                    self.dataset_admin, self.source_env_admin,
+                    SessionHelper.get_delegation_role_arn(self.source_account_id)
+                ]
             )]
             admin_statement = {
                 "Sid": DATAALL_ALLOW_OWNER_SID,
@@ -324,29 +330,41 @@ def update_dataset_bucket_key_policy(self):
             'Updating dataset Bucket KMS key policy...'
         )
         key_alias = f"alias/{self.dataset.KmsAlias}"
-        kms_client = KmsClient(account_id=self.source_account_id, region=self.source_environment.region)
+        kms_client = KmsClient(self.source_account_id, self.source_environment.region)
         kms_key_id = kms_client.get_key_id(key_alias)
         existing_policy = kms_client.get_key_policy(kms_key_id)
-        target_requester_id = SessionHelper.get_role_id(self.target_account_id, self.target_requester_IAMRoleName)
-        if existing_policy and f'{target_requester_id}:*' not in existing_policy:
-            policy = json.loads(existing_policy)
-            policy["Statement"].append(
-                {
-                    "Sid": f"{target_requester_id}",
-                    "Effect": "Allow",
-                    "Principal": {
-                        "AWS": "*"
-                    },
-                    "Action": "kms:Decrypt",
-                    "Resource": "*",
-                    "Condition": {
-                        "StringLike": {
-                            "aws:userId": f"{target_requester_id}:*"
-                        }
-                    }
-                }
-            )
-            kms_client.put_key_policy(kms_key_id, json.dumps(policy))
+        target_requester_arn = self.get_role_arn(self.target_account_id, self.target_requester_IAMRoleName)
+
+        if existing_policy:
+            existing_policy = json.loads(existing_policy)
+            counter = count()
+            statements = {item.get("Sid", next(counter)): item for item in existing_policy.get("Statement", {})}
+            if DATAALL_ACCESS_POINT_KMS_DECRYPT_SID in statements.keys():
+                logger.info(
+                    f'KMS key policy contains share statement {DATAALL_ACCESS_POINT_KMS_DECRYPT_SID}, '
+                    f'updating the current one')
+                statements[DATAALL_ACCESS_POINT_KMS_DECRYPT_SID] = (self.add_target_arn_to_statement_principal
+                                                                    (statements[DATAALL_ACCESS_POINT_KMS_DECRYPT_SID],
+                                                                     target_requester_arn))
+            else:
+                logger.info(
+                    f'KMS key does not contain share statement {DATAALL_ACCESS_POINT_KMS_DECRYPT_SID}, '
+                    f'generating a new one')
+                statements[DATAALL_ACCESS_POINT_KMS_DECRYPT_SID] = (self.generate_default_kms_decrypt_policy_statement
+                                                                    (target_requester_arn))
+            existing_policy["Statement"] = list(statements.values())
+        else:
+            logger.info('KMS key policy does not contain any statements, generating a new one')
+            existing_policy = {
+                "Version": "2012-10-17",
+                "Statement": [
+                    self.generate_default_kms_decrypt_policy_statement(target_requester_arn)
+                ]
+            }
+        kms_client.put_key_policy(
+            kms_key_id,
+            json.dumps(existing_policy)
+        )
 
     def delete_access_point_policy(self):
         logger.info(
@@ -447,24 +465,33 @@ def delete_target_role_access_policy(
                     json.dumps(existing_policy),
                 )
 
-    @staticmethod
     def delete_dataset_bucket_key_policy(
-            share: ShareObject,
+            self,
             dataset: Dataset,
-            target_environment: Environment,
     ):
         logger.info(
             'Deleting dataset bucket KMS key policy...'
         )
         key_alias = f"alias/{dataset.KmsAlias}"
-        kms_client = KmsClient(account_id=dataset.AwsAccountId, region=dataset.region)
+        kms_client = KmsClient(dataset.AwsAccountId, dataset.region)
         kms_key_id = kms_client.get_key_id(key_alias)
-        existing_policy = kms_client.get_key_policy(kms_key_id)
-        target_requester_id = SessionHelper.get_role_id(target_environment.AwsAccountId, share.principalIAMRoleName)
-        if existing_policy and f'{target_requester_id}:*' in existing_policy:
-            policy = json.loads(existing_policy)
-            policy["Statement"] = [item for item in policy["Statement"] if item.get("Sid", None) != f"{target_requester_id}"]
-            kms_client.put_key_policy(kms_key_id, json.dumps(policy))
+        existing_policy = json.loads(kms_client.get_key_policy(kms_key_id))
+        target_requester_arn = self.get_role_arn(self.target_account_id, self.target_requester_IAMRoleName)
+        counter = count()
+        statements = {item.get("Sid", next(counter)): item for item in existing_policy.get("Statement", {})}
+        if DATAALL_ACCESS_POINT_KMS_DECRYPT_SID in statements.keys():
+            principal_list = self.get_principal_list(statements[DATAALL_ACCESS_POINT_KMS_DECRYPT_SID])
+            if f"{target_requester_arn}" in principal_list:
+                principal_list.remove(f"{target_requester_arn}")
+                if len(principal_list) == 0:
+                    statements.pop(DATAALL_ACCESS_POINT_KMS_DECRYPT_SID)
+                else:
+                    statements[DATAALL_ACCESS_POINT_KMS_DECRYPT_SID]["Principal"]["AWS"] = principal_list
+                existing_policy["Statement"] = list(statements.values())
+                kms_client.put_key_policy(
+                    kms_key_id,
+                    json.dumps(existing_policy)
+                )
 
     def handle_share_failure(self, error: Exception) -> None:
         """
@@ -499,3 +526,35 @@ def handle_revoke_failure(self, error: Exception) -> None:
         DatasetAlarmService().trigger_revoke_folder_sharing_failure_alarm(
             self.target_folder, self.share, self.target_environment
         )
+
+    @staticmethod
+    def get_role_arn(target_account_id, target_requester_IAMRoleName):
+        return f"arn:aws:iam::{target_account_id}:role/{target_requester_IAMRoleName}"
+
+    @staticmethod
+    def generate_default_kms_decrypt_policy_statement(target_requester_arn):
+        return {
+            "Sid": f"{DATAALL_ACCESS_POINT_KMS_DECRYPT_SID}",
+            "Effect": "Allow",
+            "Principal": {
+                "AWS": [
+                    f"{target_requester_arn}"
+                ]
+            },
+            "Action": "kms:Decrypt",
+            "Resource": "*"
+        }
+
+    def add_target_arn_to_statement_principal(self, statement, target_requester_arn):
+        principal_list = self.get_principal_list(statement)
+        if f"{target_requester_arn}" not in principal_list:
+            principal_list.append(f"{target_requester_arn}")
+        statement["Principal"]["AWS"] = principal_list
+        return statement
+
+    @staticmethod
+    def get_principal_list(statement):
+        principal_list = statement["Principal"]["AWS"]
+        if isinstance(principal_list, str):
+            principal_list = [principal_list]
+        return principal_list

backend/dataall/modules/dataset_sharing/services/share_managers/s3_bucket_share_manager.py

@@ -19,6 +19,7 @@
 DATAALL_READ_ONLY_SID = "DataAll-Bucket-ReadOnly"
 DATAALL_ALLOW_OWNER_SID = "AllowAllToAdmin"
 IAM_S3BUCKET_ROLE_POLICY = "dataall-targetDatasetS3Bucket-AccessControlPolicy"
+DATAALL_BUCKET_KMS_DECRYPT_SID = "DataAll-Bucket-KMS-Decrypt"
 
 
 class S3BucketShareManager:
@@ -268,29 +269,33 @@ def grant_dataset_bucket_key_policy(self):
             kms_client = KmsClient(self.source_account_id, self.source_environment.region)
             kms_key_id = kms_client.get_key_id(key_alias)
             existing_policy = kms_client.get_key_policy(kms_key_id)
-            target_requester_id = SessionHelper.get_role_id(self.target_account_id, self.target_requester_IAMRoleName)
-            if existing_policy and f'{target_requester_id}:*' not in existing_policy:
-                policy = json.loads(existing_policy)
-                policy["Statement"].append(
-                    {
-                        "Sid": f"{target_requester_id}",
-                        "Effect": "Allow",
-                        "Principal": {
-                            "AWS": "*"
-                        },
-                        "Action": "kms:Decrypt",
-                        "Resource": "*",
-                        "Condition": {
-                            "StringLike": {
-                                "aws:userId": f"{target_requester_id}:*"
-                            }
-                        }
-                    }
-                )
-                kms_client.put_key_policy(
-                    kms_key_id,
-                    json.dumps(policy)
-                )
+            target_requester_arn = self.get_role_arn(self.target_account_id, self.target_requester_IAMRoleName)
+            if existing_policy:
+                existing_policy = json.loads(existing_policy)
+                counter = count()
+                statements = {item.get("Sid", next(counter)): item for item in existing_policy.get("Statement", {})}
+                if DATAALL_BUCKET_KMS_DECRYPT_SID in statements.keys():
+                    logger.info(f'KMS key policy contains share statement {DATAALL_BUCKET_KMS_DECRYPT_SID}, updating the current one')
+                    statements[DATAALL_BUCKET_KMS_DECRYPT_SID] = self.add_target_arn_to_statement_principal(
+                        statements[DATAALL_BUCKET_KMS_DECRYPT_SID], target_requester_arn)
+                else:
+                    logger.info(
+                        f'KMS key does not contain share statement {DATAALL_BUCKET_KMS_DECRYPT_SID}, generating a new one')
+                    statements[DATAALL_BUCKET_KMS_DECRYPT_SID] = self.generate_default_kms_decrypt_policy_statement(
+                        target_requester_arn)
+                existing_policy["Statement"] = list(statements.values())
+            else:
+                logger.info('KMS key policy does not contain any statements, generating a new one')
+                existing_policy = {
+                    "Version": "2012-10-17",
+                    "Statement": [
+                        self.generate_default_kms_decrypt_policy_statement(target_requester_arn)
+                    ]
+                }
+            kms_client.put_key_policy(
+                kms_key_id,
+                json.dumps(existing_policy)
+            )
 
     def delete_target_role_bucket_policy(self):
         logger.info(
@@ -374,11 +379,9 @@ def delete_target_role_access_policy(
                     json.dumps(existing_policy),
                 )
 
-    @staticmethod
     def delete_target_role_bucket_key_policy(
-            share: ShareObject,
+            self,
             target_bucket: DatasetBucket,
-            target_environment: Environment,
     ):
         if (target_bucket.imported and target_bucket.importedKmsKey) or not target_bucket.imported:
             logger.info(
@@ -387,15 +390,23 @@ def delete_target_role_bucket_key_policy(
             key_alias = f"alias/{target_bucket.KmsAlias}"
             kms_client = KmsClient(target_bucket.AwsAccountId, target_bucket.region)
             kms_key_id = kms_client.get_key_id(key_alias)
-            existing_policy = kms_client.get_key_policy(kms_key_id)
-            target_requester_id = SessionHelper.get_role_id(target_environment.AwsAccountId, share.principalIAMRoleName)
-            if existing_policy and f'{target_requester_id}:*' in existing_policy:
-                policy = json.loads(existing_policy)
-                policy["Statement"] = [item for item in policy["Statement"] if item.get("Sid", None) != f"{target_requester_id}"]
-                kms_client.put_key_policy(
-                    kms_key_id,
-                    json.dumps(policy)
-                )
+            existing_policy = json.loads(kms_client.get_key_policy(kms_key_id))
+            target_requester_arn = self.get_role_arn(self.target_account_id, self.target_requester_IAMRoleName)
+            counter = count()
+            statements = {item.get("Sid", next(counter)): item for item in existing_policy.get("Statement", {})}
+            if DATAALL_BUCKET_KMS_DECRYPT_SID in statements.keys():
+                principal_list = self.get_principal_list(statements[DATAALL_BUCKET_KMS_DECRYPT_SID])
+                if f"{target_requester_arn}" in principal_list:
+                    principal_list.remove(f"{target_requester_arn}")
+                    if len(principal_list) == 0:
+                        statements.pop(DATAALL_BUCKET_KMS_DECRYPT_SID)
+                    else:
+                        statements[DATAALL_BUCKET_KMS_DECRYPT_SID]["Principal"]["AWS"] = principal_list
+                    existing_policy["Statement"] = list(statements.values())
+                    kms_client.put_key_policy(
+                        kms_key_id,
+                        json.dumps(existing_policy)
+                    )
 
     def handle_share_failure(self, error: Exception) -> bool:
         """
@@ -456,3 +467,17 @@ def generate_default_bucket_read_policy_statement(s3_bucket_name, target_request
                 f"arn:aws:s3:::{s3_bucket_name}/*"
             ]
         }
+
+    @staticmethod
+    def generate_default_kms_decrypt_policy_statement(target_requester_arn):
+        return {
+            "Sid": f"{DATAALL_BUCKET_KMS_DECRYPT_SID}",
+            "Effect": "Allow",
+            "Principal": {
+                "AWS": [
+                    f"{target_requester_arn}"
+                ]
+            },
+            "Action": "kms:Decrypt",
+            "Resource": "*"
+        }

backend/dataall/modules/dataset_sharing/services/share_processors/s3_access_point_process_share.py

@@ -90,7 +90,8 @@ def process_approved_shares(
                 sharing_folder.manage_bucket_policy()
                 sharing_folder.grant_target_role_access_policy()
                 sharing_folder.manage_access_point_and_policy()
-                sharing_folder.update_dataset_bucket_key_policy()
+                if not dataset.imported or dataset.importedKmsKey:
+                    sharing_folder.update_dataset_bucket_key_policy()
 
                 new_state = shared_item_SM.run_transition(ShareItemActions.Success.value)
                 shared_item_SM.update_state_single_item(session, sharing_item, new_state)
@@ -177,7 +178,6 @@ def clean_up_share(
             target_environment: Environment,
             source_env_group: EnvironmentGroup,
             env_group: EnvironmentGroup,
-            existing_shared_buckets: bool = False
     ):
         """
         1) deletes S3 access point for this share in this Dataset S3 Bucket
@@ -209,11 +209,7 @@ def clean_up_share(
                 dataset=dataset,
                 target_environment=target_environment
             )
-            if not existing_shared_buckets:
-                clean_up_folder.delete_dataset_bucket_key_policy(
-                    share=share,
-                    dataset=dataset,
-                    target_environment=target_environment
-                )
+            if not dataset.imported or dataset.importedKmsKey:
+                clean_up_folder.delete_dataset_bucket_key_policy(dataset=dataset)
 
         return True

backend/dataall/modules/dataset_sharing/services/share_processors/s3_bucket_process_share.py

@@ -85,7 +85,8 @@ def process_approved_shares(
             try:
                 sharing_bucket.grant_role_bucket_policy()
                 sharing_bucket.grant_s3_iam_access()
-                sharing_bucket.grant_dataset_bucket_key_policy()
+                if not dataset.imported or dataset.importedKmsKey:
+                    sharing_bucket.grant_dataset_bucket_key_policy()
                 new_state = shared_item_SM.run_transition(ShareItemActions.Success.value)
                 shared_item_SM.update_state_single_item(session, sharing_item, new_state)
 
@@ -107,7 +108,6 @@ def process_revoked_shares(
             target_environment: Environment,
             source_env_group: EnvironmentGroup,
             env_group: EnvironmentGroup,
-            existing_shared_folders: bool = False
     ) -> bool:
         """
         1) update_share_item_status with Start action
@@ -153,11 +153,9 @@ def process_revoked_shares(
                     target_bucket=revoked_bucket,
                     target_environment=target_environment
                 )
-                if not existing_shared_folders:
+                if not dataset.imported or dataset.importedKmsKey:
                     removing_bucket.delete_target_role_bucket_key_policy(
-                        share=share,
                         target_bucket=revoked_bucket,
-                        target_environment=target_environment
                     )
                 new_state = revoked_item_SM.run_transition(ShareItemActions.Success.value)
                 revoked_item_SM.update_state_single_item(session, removing_item, new_state)