Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fixing s3 bucket sharing for federated roles #920

Merged
merged 2 commits into from
Dec 18, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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 {
Expand Down
34 changes: 16 additions & 18 deletions tests/modules/datasets/tasks/test_s3_access_point_share_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):

Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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 = {
Expand All @@ -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": "*"
Expand Down Expand Up @@ -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 = {
Expand All @@ -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": "*"
Expand Down
26 changes: 26 additions & 0 deletions tests/modules/datasets/tasks/test_s3_bucket_share_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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",
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
Loading