Skip to content

Commit

Permalink
bugfix: Fixing KMS policy explosion: Addressing PR comments
Browse files Browse the repository at this point in the history
  • Loading branch information
@anushka-singh
anushka-singh committed Nov 21, 2023
1 parent 7665a5e commit bb34b24
Show file tree
Hide file tree
Showing 6 changed files with 64 additions and 39 deletions.
6 changes: 2 additions & 4 deletions backend/dataall/modules/dataset_sharing/services/data_sharing_service.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}")

Expand All @@ -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}')

Expand Down
35 changes: 22 additions & 13 deletions .../dataall/modules/dataset_sharing/services/share_managers/s3_access_point_share_manager.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -337,20 +337,29 @@ def update_dataset_bucket_key_policy(self):

counter = count()
statements = {item.get("Sid", next(counter)): item for item in existing_policy.get("Statement", {})}
if existing_policy and 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))
if existing_policy:
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(
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())
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)
Expand Down
28 changes: 18 additions & 10 deletions backend/dataall/modules/dataset_sharing/services/share_managers/s3_bucket_share_manager.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -272,17 +272,25 @@ def grant_dataset_bucket_key_policy(self):
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 existing_policy and 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)
if existing_policy:
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(
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())
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)
Expand Down
8 changes: 3 additions & 5 deletions ...ataall/modules/dataset_sharing/services/share_processors/s3_access_point_process_share.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -177,7 +177,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
Expand Down Expand Up @@ -209,9 +208,8 @@ def clean_up_share(
dataset=dataset,
target_environment=target_environment
)
if not existing_shared_buckets:
clean_up_folder.delete_dataset_bucket_key_policy(
dataset=dataset
)
clean_up_folder.delete_dataset_bucket_key_policy(
dataset=dataset
)

return True
8 changes: 3 additions & 5 deletions backend/dataall/modules/dataset_sharing/services/share_processors/s3_bucket_process_share.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,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
Expand Down Expand Up @@ -153,10 +152,9 @@ def process_revoked_shares(
target_bucket=revoked_bucket,
target_environment=target_environment
)
if not existing_shared_folders:
removing_bucket.delete_target_role_bucket_key_policy(
target_bucket=revoked_bucket,
)
removing_bucket.delete_target_role_bucket_key_policy(
target_bucket=revoked_bucket,
)
new_state = revoked_item_SM.run_transition(ShareItemActions.Success.value)
revoked_item_SM.update_state_single_item(session, removing_item, new_state)

Expand Down
18 changes: 16 additions & 2 deletions tests/modules/datasets/tasks/test_s3_access_point_share_manager.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1264,7 +1264,14 @@ 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"
target_requester_arn = f"arn:aws:iam::{TARGET_ACCOUNT_ENV}:role/{TARGET_ACCOUNT_ENV_ROLE_NAME}"
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
)

# Includes target env admin to be removed and another, that should remain
existing_key_policy = {
Expand Down Expand Up @@ -1339,7 +1346,14 @@ 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"
target_requester_arn = f"arn:aws:iam::{TARGET_ACCOUNT_ENV}:role/{TARGET_ACCOUNT_ENV_ROLE_NAME}"
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
)

# Includes target env admin to be removed and another, that should remain
existing_key_policy = {
Expand Down

0 comments on commit bb34b24

Please sign in to comment.