From 3309ff5352f468071d922d596d2501b0ef2f9335 Mon Sep 17 00:00:00 2001 From: FastLee Date: Wed, 17 Apr 2024 09:35:46 -0400 Subject: [PATCH] Fixed update UC Trust Role to amend existing statement, rather than rewriting all statements. --- src/databricks/labs/ucx/assessment/aws.py | 67 ++++++++++++----------- tests/unit/assessment/test_aws.py | 55 +++++++++++++++++-- 2 files changed, 84 insertions(+), 38 deletions(-) diff --git a/src/databricks/labs/ucx/assessment/aws.py b/src/databricks/labs/ucx/assessment/aws.py index 97276dc701..0c318e3922 100644 --- a/src/databricks/labs/ucx/assessment/aws.py +++ b/src/databricks/labs/ucx/assessment/aws.py @@ -216,18 +216,20 @@ def _s3_actions(self, actions): return s3_actions def _aws_role_trust_doc(self, external_id="0000"): + return self._get_json_for_cli( + { + "Version": "2012-10-17", + "Statement": [self._databricks_trust_statement(external_id)], + } + ) + + @staticmethod + def _databricks_trust_statement(external_id="0000"): return { - "Version": "2012-10-17", - "Statement": [ - { - "Effect": "Allow", - "Principal": { - "AWS": "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL" - }, - "Action": "sts:AssumeRole", - "Condition": {"StringEquals": {"sts:ExternalId": external_id}}, - } - ], + "Effect": "Allow", + "Principal": {"AWS": "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL"}, + "Action": "sts:AssumeRole", + "Condition": {"StringEquals": {"sts:ExternalId": external_id}}, } def _aws_s3_policy(self, s3_prefixes, account_id, role_name, kms_key=None): @@ -285,7 +287,7 @@ def create_uc_role(self, role_name: str) -> str | None: the AssumeRole condition will be modified later with the external ID captured from the UC credential. https://docs.databricks.com/en/connect/unity-catalog/storage-credentials.html """ - return self._create_role(role_name, self._get_json_for_cli(self._aws_role_trust_doc())) + return self._create_role(role_name, self._aws_role_trust_doc()) def update_uc_trust_role(self, role_name: str, external_id: str = "0000") -> str | None: """ @@ -299,27 +301,28 @@ def update_uc_trust_role(self, role_name: str, external_id: str = "0000") -> str logger.error(f"Role {role_name} doesn't exist") return None policy_document = role.get("AssumeRolePolicyDocument") - if not policy_document: - logger.error(f"Role {role_name} doesn't have an AssumeRolePolicyDocument") - return None - for idx, statement in enumerate(policy_document["Statement"]): - effect = statement.get("Effect") - action = statement.get("Action") - principal = statement.get("Principal") - if not (effect and action and principal): - continue - if effect != "Allow": - continue - if action != "sts:AssumeRole": - continue - principal = principal.get("AWS") - if not principal: - continue - if not self._is_uc_principal(principal): - continue - policy_document["Statement"][idx] = self._aws_role_trust_doc(external_id) + if policy_document and policy_document.get("Statement"): + for idx, statement in enumerate(policy_document["Statement"]): + effect = statement.get("Effect") + action = statement.get("Action") + principal = statement.get("Principal") + if not (effect and action and principal): + continue + if effect != "Allow": + continue + if action != "sts:AssumeRole": + continue + principal = principal.get("AWS") + if not principal: + continue + if not self._is_uc_principal(principal): + continue + policy_document["Statement"][idx] = self._databricks_trust_statement(external_id) + policy_document_json = self._get_json_for_cli(policy_document) + else: + policy_document_json = self._aws_role_trust_doc(external_id) update_role = self._run_json_command( - f"iam update-assume-role-policy --role-name {role_name} --policy-document {self._get_json_for_cli(policy_document)}" + f"iam update-assume-role-policy --role-name {role_name} " f"--policy-document {policy_document_json}" ) if not update_role: return None diff --git a/tests/unit/assessment/test_aws.py b/tests/unit/assessment/test_aws.py index dd15d20b44..ffc5678813 100644 --- a/tests/unit/assessment/test_aws.py +++ b/tests/unit/assessment/test_aws.py @@ -540,7 +540,7 @@ def command_call(cmd: str): ) in command_calls -def test_update_uc_trust_role(mocker): +def test_update_uc_trust_role_append(mocker): command_calls = [] mocker.patch("shutil.which", return_value="/path/aws") @@ -592,11 +592,54 @@ def command_call(cmd: str): assert ( '/path/aws iam update-assume-role-policy --role-name test_role ' '--policy-document ' - '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":' - '{"AWS":"arn:aws:iam::0123456789:root"},"Action":"sts:AssumeRole"},' - '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":' - '{"AWS":"arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL"}' - ',"Action":"sts:AssumeRole","Condition":{"StringEquals":{"sts:ExternalId":"1234"}}}]}]} ' + '{"Version":"2012-10-17",' + '"Statement":[' + '{"Effect":"Allow","Principal":{"AWS":"arn:aws:iam::0123456789:root"},' + '"Action":"sts:AssumeRole"},' + '{"Effect":"Allow",' + '"Principal":{"AWS":"arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL"},' + '"Action":"sts:AssumeRole","Condition":{"StringEquals":{"sts:ExternalId":"1234"}}}]} ' + '--output json' + ) in command_calls + + +def test_update_uc_trust_role(mocker): + command_calls = [] + mocker.patch("shutil.which", return_value="/path/aws") + + def command_call(cmd: str): + if "iam get-role" in cmd: + return ( + 0, + """ +{ + "Role": { + "Path": "/", + "RoleName": "Test-Role", + "RoleId": "ABCD", + "Arn": "arn:aws:iam::0123456789:role/Test-Role", + "CreateDate": "2024-01-01T12:00:00+00:00", + "Description": "", + "MaxSessionDuration": 3600, + "RoleLastUsed": {} + } +} + """, + "", + ) + command_calls.append(cmd) + return 0, '{"Role": {"Arn": "arn:aws:iam::123456789012:role/Test-Role"}}', "" + + aws = AWSResources("Fake_Profile", command_call) + aws.update_uc_trust_role("test_role", "1234") + assert ( + '/path/aws iam update-assume-role-policy --role-name test_role ' + '--policy-document ' + '{"Version":"2012-10-17",' + '"Statement":[' + '{"Effect":"Allow",' + '"Principal":{"AWS":"arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL"},' + '"Action":"sts:AssumeRole","Condition":{"StringEquals":{"sts:ExternalId":"1234"}}}]} ' '--output json' ) in command_calls