From d937db3ac470d3bdb9854acdd9ac428d7d0041bc Mon Sep 17 00:00:00 2001 From: Zilvinas Saltys Date: Fri, 15 Dec 2023 17:02:24 +0000 Subject: [PATCH 1/2] fixing s3 bucket sharing for federated roles --- .../share_managers/s3_bucket_share_manager.py | 27 ++++++++++--------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/backend/dataall/modules/dataset_sharing/services/share_managers/s3_bucket_share_manager.py b/backend/dataall/modules/dataset_sharing/services/share_managers/s3_bucket_share_manager.py index 308aa5e7d..375c1b94c 100644 --- a/backend/dataall/modules/dataset_sharing/services/share_managers/s3_bucket_share_manager.py +++ b/backend/dataall/modules/dataset_sharing/services/share_managers/s3_bucket_share_manager.py @@ -204,19 +204,23 @@ def grant_role_bucket_policy(self): f'Granting access via Bucket policy for {self.bucket_name}' ) try: - target_requester_arn = self.get_role_arn(self.target_account_id, self.target_requester_IAMRoleName) + target_requester_arn = IAM.get_role_arn_by_name(self.target_account_id, self.target_requester_IAMRoleName) bucket_policy = self.get_bucket_policy_or_default() counter = count() statements = {item.get("Sid", next(counter)): item for item in bucket_policy.get("Statement", {})} if DATAALL_READ_ONLY_SID in statements.keys(): logger.info(f'Bucket policy contains share statement {DATAALL_READ_ONLY_SID}, updating the current one') - statements[DATAALL_READ_ONLY_SID] = self.add_target_arn_to_statement_principal(statements[DATAALL_READ_ONLY_SID], target_requester_arn) + statements[DATAALL_READ_ONLY_SID] = self.add_target_arn_to_statement_principal( + statements[DATAALL_READ_ONLY_SID], target_requester_arn) else: - logger.info(f'Bucket policy does not contain share statement {DATAALL_READ_ONLY_SID}, generating a new one') - statements[DATAALL_READ_ONLY_SID] = self.generate_default_bucket_read_policy_statement(self.bucket_name, target_requester_arn) + logger.info( + f'Bucket policy does not contain share statement {DATAALL_READ_ONLY_SID}, generating a new one') + statements[DATAALL_READ_ONLY_SID] = self.generate_default_bucket_read_policy_statement(self.bucket_name, + target_requester_arn) if DATAALL_ALLOW_OWNER_SID not in statements.keys(): - statements[DATAALL_ALLOW_OWNER_SID] = self.generate_owner_access_statement(self.bucket_name, self.get_bucket_owner_roleid()) + statements[DATAALL_ALLOW_OWNER_SID] = self.generate_owner_access_statement(self.bucket_name, + self.get_bucket_owner_roleid()) bucket_policy["Statement"] = list(statements.values()) s3_client = S3Client(self.source_account_id, self.source_environment.region) @@ -269,13 +273,14 @@ 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_arn = self.get_role_arn(self.target_account_id, self.target_requester_IAMRoleName) + target_requester_arn = IAM.get_role_arn_by_name(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') + 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: @@ -304,7 +309,7 @@ def delete_target_role_bucket_policy(self): try: s3_client = S3Client(self.source_account_id, self.source_environment.region) bucket_policy = json.loads(s3_client.get_bucket_policy(self.bucket_name)) - target_requester_arn = self.get_role_arn(self.target_account_id, self.target_requester_IAMRoleName) + target_requester_arn = IAM.get_role_arn_by_name(self.target_account_id, self.target_requester_IAMRoleName) counter = count() statements = {item.get("Sid", next(counter)): item for item in bucket_policy.get("Statement", {})} if DATAALL_READ_ONLY_SID in statements.keys(): @@ -391,7 +396,7 @@ def delete_target_role_bucket_key_policy( kms_client = KmsClient(target_bucket.AwsAccountId, target_bucket.region) kms_key_id = kms_client.get_key_id(key_alias) 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) + target_requester_arn = IAM.get_role_arn_by_name(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(): @@ -444,10 +449,6 @@ def handle_revoke_failure(self, error: Exception) -> bool: ) return True - @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_bucket_read_policy_statement(s3_bucket_name, target_requester_arn): return { From b61d02fa27f3f3cb478a22aff613f5405dcf88f8 Mon Sep 17 00:00:00 2001 From: Noah Paige Date: Fri, 15 Dec 2023 13:36:35 -0500 Subject: [PATCH 2/2] Fix integration tests and get IAM arn in folder share --- .../s3_access_point_share_manager.py | 8 ++--- .../test_s3_access_point_share_manager.py | 34 +++++++++---------- .../tasks/test_s3_bucket_share_manager.py | 26 ++++++++++++++ 3 files changed, 44 insertions(+), 24 deletions(-) diff --git a/backend/dataall/modules/dataset_sharing/services/share_managers/s3_access_point_share_manager.py b/backend/dataall/modules/dataset_sharing/services/share_managers/s3_access_point_share_manager.py index 610657df7..ae1c02ca4 100644 --- a/backend/dataall/modules/dataset_sharing/services/share_managers/s3_access_point_share_manager.py +++ b/backend/dataall/modules/dataset_sharing/services/share_managers/s3_access_point_share_manager.py @@ -333,7 +333,7 @@ def update_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_arn = self.get_role_arn(self.target_account_id, self.target_requester_IAMRoleName) + target_requester_arn = IAM.get_role_arn_by_name(self.target_account_id, self.target_requester_IAMRoleName) if existing_policy: existing_policy = json.loads(existing_policy) @@ -476,7 +476,7 @@ def delete_dataset_bucket_key_policy( kms_client = KmsClient(dataset.AwsAccountId, dataset.region) kms_key_id = kms_client.get_key_id(key_alias) 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) + target_requester_arn = IAM.get_role_arn_by_name(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(): @@ -527,10 +527,6 @@ def handle_revoke_failure(self, error: Exception) -> None: 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 { diff --git a/tests/modules/datasets/tasks/test_s3_access_point_share_manager.py b/tests/modules/datasets/tasks/test_s3_access_point_share_manager.py index cc175f179..b111d9daa 100644 --- a/tests/modules/datasets/tasks/test_s3_access_point_share_manager.py +++ b/tests/modules/datasets/tasks/test_s3_access_point_share_manager.py @@ -180,6 +180,16 @@ def mock_kms_client(mocker): return mock_client +def mock_iam_client(mocker, account_id, role_name): + mock_client = MagicMock() + mocker.patch( + 'dataall.modules.dataset_sharing.services.share_managers.s3_access_point_share_manager.IAM', + mock_client + ) + mock_client.get_role_arn_by_name.return_value = f"arn:aws:iam::{account_id}:role/{role_name}" + return mock_client + + @pytest.fixture(scope="module") def target_dataset_access_control_policy(request): @@ -508,6 +518,7 @@ def test_update_dataset_bucket_key_policy_with_env_admin( # Given kms_client = mock_kms_client(mocker) kms_client().get_key_id.return_value = None + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share1.principalIAMRoleName) existing_key_policy = { "Version": "2012-10-17", @@ -613,6 +624,7 @@ def test_update_dataset_bucket_key_policy_without_env_admin( # Given kms_client = mock_kms_client(mocker) kms_client().get_key_id.return_value = "kms-key" + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share1.principalIAMRoleName) existing_key_policy = { "Version": "2012-10-17", @@ -1264,14 +1276,7 @@ def test_delete_dataset_bucket_key_policy_existing_policy_with_additional_target # Given kms_client = mock_kms_client(mocker) kms_client().get_key_id.return_value = "1" - s3_access_point_share_manager_mocker = S3AccessPointShareManager(mocker.MagicMock(), mocker.MagicMock(), - mocker.MagicMock(), mocker.MagicMock(), - mocker.MagicMock(), mocker.MagicMock(), - mocker.MagicMock(), mocker.MagicMock()) - target_requester_arn = s3_access_point_share_manager_mocker.get_role_arn( - TARGET_ACCOUNT_ENV, - TARGET_ACCOUNT_ENV_ROLE_NAME - ) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share1.principalIAMRoleName) # Includes target env admin to be removed and another, that should remain existing_key_policy = { @@ -1282,7 +1287,7 @@ def test_delete_dataset_bucket_key_policy_existing_policy_with_additional_target "Effect": "Allow", "Principal": {"AWS": [ "SomeTargetResourceArn", - f"{target_requester_arn}" + f"arn:aws:iam::{target_environment.AwsAccountId}:role/{share1.principalIAMRoleName}" ]}, "Action": "kms:Decrypt", "Resource": "*" @@ -1346,14 +1351,7 @@ def test_delete_dataset_bucket_key_policy_existing_policy_with_no_additional_tar # Given kms_client = mock_kms_client(mocker) kms_client().get_key_id.return_value = "1" - s3_access_point_share_manager_mocker = S3AccessPointShareManager(mocker.MagicMock(), mocker.MagicMock(), - mocker.MagicMock(), mocker.MagicMock(), - mocker.MagicMock(), mocker.MagicMock(), - mocker.MagicMock(), mocker.MagicMock()) - target_requester_arn = s3_access_point_share_manager_mocker.get_role_arn( - TARGET_ACCOUNT_ENV, - TARGET_ACCOUNT_ENV_ROLE_NAME - ) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share1.principalIAMRoleName) # Includes target env admin to be removed and another, that should remain existing_key_policy = { @@ -1363,7 +1361,7 @@ def test_delete_dataset_bucket_key_policy_existing_policy_with_no_additional_tar "Sid": f"{DATAALL_ACCESS_POINT_KMS_DECRYPT_SID}", "Effect": "Allow", "Principal": {"AWS": [ - f"{target_requester_arn}" + f"arn:aws:iam::{target_environment.AwsAccountId}:role/{share1.principalIAMRoleName}" ]}, "Action": "kms:Decrypt", "Resource": "*" diff --git a/tests/modules/datasets/tasks/test_s3_bucket_share_manager.py b/tests/modules/datasets/tasks/test_s3_bucket_share_manager.py index c11a1ef36..01e4cf691 100644 --- a/tests/modules/datasets/tasks/test_s3_bucket_share_manager.py +++ b/tests/modules/datasets/tasks/test_s3_bucket_share_manager.py @@ -216,6 +216,16 @@ def mock_kms_client(mocker): return mock_client +def mock_iam_client(mocker, account_id, role_name): + mock_client = MagicMock() + mocker.patch( + 'dataall.modules.dataset_sharing.services.share_managers.s3_bucket_share_manager.IAM', + mock_client + ) + mock_client.get_role_arn_by_name.return_value = f"arn:aws:iam::{account_id}:role/{role_name}" + return mock_client + + # For below test cases, dataset2, share2, src, target env and src group , env group remain the same def test_grant_role_bucket_policy_with_no_policy_present( mocker, @@ -232,6 +242,7 @@ def test_grant_role_bucket_policy_with_no_policy_present( # No Bucket policy. A Default bucket policy should be formed with DataAll-Bucket-ReadOnly, AllowAllToAdmin & RequiredSecureTransport Sids s3_client = mock_s3_client(mocker) s3_client().get_bucket_policy.return_value = None + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) mocker.patch( "dataall.base.aws.sts.SessionHelper.get_delegation_role_arn", @@ -299,6 +310,7 @@ def test_grant_role_bucket_policy_with_default_complete_policy( s3_client = mock_s3_client(mocker) s3_client().get_bucket_policy.return_value = json.dumps(bucket_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager( @@ -344,6 +356,7 @@ def test_grant_role_bucket_policy_with_policy_and_no_allow_owner_sid_and_no_read s3_client = mock_s3_client(mocker) s3_client().get_bucket_policy.return_value = json.dumps(bucket_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) mocker.patch( "dataall.base.aws.sts.SessionHelper.get_delegation_role_arn", @@ -419,6 +432,7 @@ def test_grant_role_bucket_policy_with_another_read_only_role( s3_client = mock_s3_client(mocker) s3_client().get_bucket_policy.return_value = json.dumps(bucket_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) mocker.patch( "dataall.base.aws.sts.SessionHelper.get_delegation_role_arn", @@ -676,6 +690,7 @@ def test_grant_dataset_bucket_key_policy_with_complete_policy_present( existing_key_policy = base_kms_key_policy() kms_client().get_key_policy.return_value = json.dumps(existing_key_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager( @@ -712,6 +727,7 @@ def test_grant_dataset_bucket_key_policy_with_target_requester_id_absent( existing_key_policy = base_kms_key_policy("OtherTargetRequestorArn") kms_client().get_key_policy.return_value = json.dumps(existing_key_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) # Mock the S3BucketShareManager with the KMS client with db.scoped_session() as session: @@ -768,6 +784,7 @@ def test_grant_dataset_bucket_key_policy_and_default_bucket_key_policy( kms_client().get_key_id.return_value = "kms-key" kms_client().get_key_policy.return_value = json.dumps(existing_key_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share3.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager( @@ -818,6 +835,7 @@ def test_grant_dataset_bucket_key_policy_with_imported( kms_client().get_key_id.return_value = "kms-key" kms_client().get_key_policy.return_value = json.dumps(existing_key_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share3.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager( @@ -864,6 +882,7 @@ def test_delete_target_role_bucket_policy_with_no_read_only_sid( s3_client = mock_s3_client(mocker) s3_client().get_bucket_policy.return_value = json.dumps(bucket_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager( @@ -923,6 +942,7 @@ def test_delete_target_role_bucket_policy_with_multiple_principals_in_policy( s3_client = mock_s3_client(mocker) s3_client().get_bucket_policy.return_value = json.dumps(bucket_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager( @@ -996,6 +1016,7 @@ def test_delete_target_role_bucket_policy_with_one_principal_in_policy( s3_client = mock_s3_client(mocker) s3_client().get_bucket_policy.return_value = json.dumps(bucket_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager( @@ -1284,6 +1305,7 @@ def test_delete_target_role_bucket_key_policy_with_no_target_requester_id( kms_client().get_key_id.return_value = "kms-key" kms_client().get_key_policy.return_value = json.dumps(existing_key_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager( @@ -1324,6 +1346,7 @@ def test_delete_target_role_bucket_key_policy_with_target_requester_id( kms_client().get_key_id.return_value = "kms-key" kms_client().get_key_policy.return_value = json.dumps(existing_key_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager( @@ -1370,6 +1393,7 @@ def test_delete_target_role_bucket_key_policy_with_target_requester_id_and_impor kms_client().get_key_id.return_value = "kms-key" kms_client().get_key_policy.return_value = json.dumps(existing_key_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share3.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager( @@ -1416,6 +1440,7 @@ def test_delete_target_role_bucket_key_policy_with_target_requester_id_and_impor kms_client().get_key_id.return_value = "kms-key" kms_client().get_key_policy.return_value = json.dumps(existing_key_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share3.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager( @@ -1479,6 +1504,7 @@ def test_delete_target_role_bucket_key_policy_with_multiple_principals_in_policy kms_client().get_key_id.return_value = "kms-key" kms_client().get_key_policy.return_value = json.dumps(existing_key_policy) + iam_client = mock_iam_client(mocker, target_environment.AwsAccountId, share2.principalIAMRoleName) with db.scoped_session() as session: manager = S3BucketShareManager(