Added AWS IAM roles support to databricks labs ucx migrate-credentials command (#973)

## Changes
<!-- Summary of your changes that are easy to understand. Add
screenshots when necessary -->
A few more things to be done
- [x] Added `load` function to `AWSResourcePermissions` to return
identified instance profiles
- [x] Added `IamRoleMigration` class under `aws/credentials.py` to
migrate AWS instance profiles identified

### Linked issues
<!-- DOC: Link issue with a keyword: close, closes, closed, fix, fixes,
fixed, resolve, resolves, resolved. See
https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword
-->

Resolves #862

Related PR:
- #874

### Functionality 

- [x] added relevant user documentation
- [x] added new CLI command `databricks labs ucx migrate-credentials`

### Tests
<!-- How is this tested? Please see the checklist below and also
describe any other relevant tests -->

- [x] manually tested
- [x] added unit tests
- [x] added integration tests

---------

Co-authored-by: qziyuan <[email protected]>

Author: nkvuong

src/databricks/labs/ucx/assessment/aws.py

@@ -44,11 +44,16 @@ class AWSRoleAction:
     privilege: str
     resource_path: str
 
+    @property
+    def role_name(self):
+        role_match = re.match(AWSInstanceProfile.ROLE_NAME_REGEX, self.role_arn)
+        return role_match.group(1)
+
 
 @dataclass
 class AWSInstanceProfile:
     instance_profile_arn: str
-    iam_role_arn: str | None = None
+    iam_role_arn: str
 
     ROLE_NAME_REGEX = r"arn:aws:iam::[0-9]+:(?:instance-profile|role)\/([a-zA-Z0-9+=,.@_-]*)$"
 
@@ -219,8 +224,8 @@ def _s3_actions(self, actions):
             s3_actions = [actions]
         return s3_actions
 
-    def add_uc_role(self, role_name):
-        aws_role_trust_doc = {
+    def _aws_role_trust_doc(self, external_id="0000"):
+        return {
             "Version": "2012-10-17",
             "Statement": [
                 {
@@ -229,21 +234,42 @@ def add_uc_role(self, role_name):
                         "AWS": "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL"
                     },
                     "Action": "sts:AssumeRole",
-                    "Condition": {"StringEquals": {"sts:ExternalId": "0000"}},
+                    "Condition": {"StringEquals": {"sts:ExternalId": external_id}},
                 }
             ],
         }
-        # the AssumeRole condition will be modified with the external ID captured from the UC credential.
-        # https://docs.databricks.com/en/connect/unity-catalog/storage-credentials.html
-        assume_role_json = self._get_json_for_cli(aws_role_trust_doc)
+
+    def add_uc_role(self, role_name: str) -> bool:
+        """
+        Create an IAM role for Unity Catalog to access the S3 buckets.
+        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
+        """
+        assume_role_json = self._get_json_for_cli(self._aws_role_trust_doc())
         add_role = self._run_json_command(
             f"iam create-role --role-name {role_name} --assume-role-policy-document {assume_role_json}"
         )
         if not add_role:
             return False
         return True
 
-    def add_uc_role_policy(self, role_name, policy_name, s3_prefixes: set[str], account_id: str, kms_key=None):
+    def update_uc_trust_role(self, role_name: str, external_id: str = "0000") -> bool:
+        """
+        Modify an existing IAM role for Unity Catalog to access the S3 buckets with the external ID
+        captured from the UC credential.
+        https://docs.databricks.com/en/connect/unity-catalog/storage-credentials.html
+        """
+        assume_role_json = self._get_json_for_cli(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 {assume_role_json}"
+        )
+        if not update_role:
+            return False
+        return True
+
+    def add_uc_role_policy(
+        self, role_name: str, policy_name: str, s3_prefixes: set[str], account_id: str, kms_key=None
+    ) -> bool:
         s3_prefixes_enriched = sorted([self.S3_PREFIX + s3_prefix for s3_prefix in s3_prefixes])
         statement = [
             {
@@ -278,7 +304,8 @@ def add_uc_role_policy(self, role_name, policy_name, s3_prefixes: set[str], acco
 
         policy_document_json = self._get_json_for_cli(policy_document)
         if not self._run_command(
-            f"iam put-role-policy --role-name {role_name} --policy-name {policy_name} --policy-document {policy_document_json}"
+            f"iam put-role-policy --role-name {role_name} "
+            f"--policy-name {policy_name} --policy-document {policy_document_json}"
         ):
             return False
         return True
@@ -325,6 +352,7 @@ def __init__(
         self._schema = schema
         self._aws_account_id = aws_account_id
         self._kms_key = kms_key
+        self._filename = self.INSTANCE_PROFILES_FILE_NAMES
 
     @classmethod
     def for_cli(cls, ws: WorkspaceClient, backend, aws_profile, schema, kms_key=None, product='ucx'):
@@ -338,26 +366,11 @@ def for_cli(cls, ws: WorkspaceClient, backend, aws_profile, schema, kms_key=None
             ws,
             backend,
             aws,
-            schema=schema,
-            aws_account_id=caller_identity.get("Account"),
-            kms_key=kms_key,
+            schema,
+            caller_identity.get("Account"),
+            kms_key,
         )
 
-    def save_uc_compatible_roles(self):
-        uc_role_access = list(self._get_role_access())
-        if len(uc_role_access) == 0:
-            logger.warning("No Mapping Was Generated.")
-            return None
-        return self._installation.save(uc_role_access, filename=self.UC_ROLES_FILE_NAMES)
-
-    def get_uc_compatible_roles(self):
-        try:
-            role_actions = self._installation.load(list[AWSRoleAction], filename=self.UC_ROLES_FILE_NAMES)
-        except ResourceDoesNotExist:
-            self.save_uc_compatible_roles()
-            role_actions = self._installation.load(list[AWSRoleAction], filename=self.UC_ROLES_FILE_NAMES)
-        return role_actions
-
     def create_uc_roles_cli(self, *, single_role=True, role_name="UC_ROLE", policy_name="UC_POLICY"):
         # Get the missing paths
         # Identify the S3 prefixes
@@ -373,7 +386,7 @@ def create_uc_roles_cli(self, *, single_role=True, role_name="UC_ROLE", policy_n
         if single_role:
             if self._aws_resources.add_uc_role(role_name):
                 self._aws_resources.add_uc_role_policy(
-                    role_name, policy_name, s3_prefixes, account_id=self._aws_account_id, kms_key=self._kms_key
+                    role_name, policy_name, s3_prefixes, self._aws_account_id, self._kms_key
                 )
         else:
             role_id = 1
@@ -383,15 +396,44 @@ def create_uc_roles_cli(self, *, single_role=True, role_name="UC_ROLE", policy_n
                         f"{role_name}-{role_id}",
                         f"{policy_name}-{role_id}",
                         {s3_prefix},
-                        account_id=self._aws_account_id,
-                        kms_key=self._kms_key,
+                        self._aws_account_id,
+                        self._kms_key,
                     )
                 role_id += 1
 
+    def update_uc_role_trust_policy(self, role_name, external_id="0000"):
+        return self._aws_resources.update_uc_trust_role(role_name, external_id)
+
+    def save_uc_compatible_roles(self):
+        uc_role_access = list(self._get_role_access())
+        if len(uc_role_access) == 0:
+            logger.warning("No mapping was generated.")
+            return None
+        return self._installation.save(uc_role_access, filename=self.UC_ROLES_FILE_NAMES)
+
+    def load_uc_compatible_roles(self):
+        try:
+            role_actions = self._installation.load(list[AWSRoleAction], filename=self.UC_ROLES_FILE_NAMES)
+        except ResourceDoesNotExist:
+            self.save_uc_compatible_roles()
+            role_actions = self._installation.load(list[AWSRoleAction], filename=self.UC_ROLES_FILE_NAMES)
+        return role_actions
+
+    def save_instance_profile_permissions(self) -> str | None:
+        instance_profile_access = list(self._get_instance_profiles_access())
+        if len(instance_profile_access) == 0:
+            logger.warning("No mapping was generated.")
+            return None
+        return self._installation.save(instance_profile_access, filename=self.INSTANCE_PROFILES_FILE_NAMES)
+
     def _get_instance_profiles(self) -> Iterable[AWSInstanceProfile]:
         instance_profiles = self._ws.instance_profiles.list()
         result_instance_profiles = []
         for instance_profile in instance_profiles:
+            if not instance_profile.iam_role_arn:
+                instance_profile.iam_role_arn = instance_profile.instance_profile_arn.replace(
+                    "instance-profile", "role"
+                )
             result_instance_profiles.append(
                 AWSInstanceProfile(instance_profile.instance_profile_arn, instance_profile.iam_role_arn)
             )
@@ -446,7 +488,7 @@ def _get_role_access_task(self, arn: str, role_name: str):
 
     def _identify_missing_paths(self):
         external_locations = ExternalLocations(self._ws, self._backend, self._schema).snapshot()
-        compatible_roles = self.get_uc_compatible_roles()
+        compatible_roles = self.load_uc_compatible_roles()
         missing_paths = set()
         for external_location in external_locations:
             path = PurePath(external_location.location)
@@ -460,13 +502,6 @@ def _identify_missing_paths(self):
             missing_paths.add(external_location.location)
         return missing_paths
 
-    def save_instance_profile_permissions(self) -> str | None:
-        instance_profile_access = list(self._get_instance_profiles_access())
-        if len(instance_profile_access) == 0:
-            logger.warning("No Mapping Was Generated.")
-            return None
-        return self._installation.save(instance_profile_access, filename=self.INSTANCE_PROFILES_FILE_NAMES)
-
     def _identify_missing_external_locations(
         self,
         external_locations: Iterable[ExternalLocation],
@@ -512,7 +547,7 @@ def create_external_locations(self, location_init="UCX_location"):
         external_locations = ExternalLocations(self._ws, self._backend, self._schema).snapshot()
         existing_external_locations = self._ws.external_locations.list()
         existing_paths = [external_location.url for external_location in existing_external_locations]
-        compatible_roles = self.get_uc_compatible_roles()
+        compatible_roles = self.load_uc_compatible_roles()
         missing_paths = self._identify_missing_external_locations(external_locations, existing_paths, compatible_roles)
         external_location_names = [external_location.name for external_location in existing_external_locations]
         external_location_num = 1