Fixed look up logic where instance profile name does not match role n…

…ame (#1716)

closes #1711
Changed instance profile lookup mechanism to actively lookup the
attached role and not rely on the role name being the same as the
instance profile name.

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

@@ -46,7 +46,7 @@ class AWSRoleAction:
 
     @property
     def role_name(self):
-        role_match = re.match(AWSInstanceProfile.ROLE_NAME_REGEX, self.role_arn)
+        role_match = re.match(AWSResources.ROLE_NAME_REGEX, self.role_arn)
         return role_match.group(1)
 
 
@@ -55,15 +55,13 @@ class AWSInstanceProfile:
     instance_profile_arn: str
     iam_role_arn: str | None = None
 
-    ROLE_NAME_REGEX = r"arn:aws:iam::[0-9]+:(?:instance-profile|role)\/([a-zA-Z0-9+=,.@_-]*)$"
-
     @property
     def role_name(self) -> str | None:
         if self.iam_role_arn:
             arn = self.iam_role_arn
         else:
             arn = self.instance_profile_arn
-        role_match = re.match(self.ROLE_NAME_REGEX, arn)
+        role_match = re.match(AWSResources.ROLE_NAME_REGEX, arn)
         if not role_match:
             logger.error(f"Role ARN is mismatched {self.iam_role_arn}")
             return None
@@ -88,6 +86,7 @@ class AWSResources:
         "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL",
         "arn:aws:iam::707343435239:role/unity-catalog-dev-UCMasterRole-G3MMN8SP21FO",
     ]
+    ROLE_NAME_REGEX = r"arn:aws:iam::[0-9]+:(?:instance-profile|role)\/([a-zA-Z0-9+=,.@_-]*)$"
 
     def __init__(self, profile: str, command_runner: Callable[[str], tuple[int, str, str]] = run_command):
         self._profile = profile
@@ -350,7 +349,7 @@ def create_migration_role(self, role_name: str) -> str | None:
         assume_role_json = self._get_json_for_cli(aws_role_trust_doc)
         return self._create_role(role_name, assume_role_json)
 
-    def get_instance_profile(self, instance_profile_name: str) -> str | None:
+    def get_instance_profile_arn(self, instance_profile_name: str) -> str | None:
         instance_profile = self._run_json_command(
             f"iam get-instance-profile --instance-profile-name {instance_profile_name}"
         )
@@ -360,6 +359,19 @@ def get_instance_profile(self, instance_profile_name: str) -> str | None:
 
         return instance_profile["InstanceProfile"]["Arn"]
 
+    def get_instance_profile_role_arn(self, instance_profile_name: str) -> str | None:
+        instance_profile = self._run_json_command(
+            f"iam get-instance-profile --instance-profile-name {instance_profile_name}"
+        )
+
+        if not instance_profile:
+            return None
+
+        try:
+            return instance_profile["InstanceProfile"]["Roles"][0]["Arn"]
+        except (KeyError, IndexError):
+            return None
+
     def create_instance_profile(self, instance_profile_name: str) -> str | None:
         instance_profile = self._run_json_command(
             f"iam create-instance-profile --instance-profile-name {instance_profile_name}"

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

@@ -111,13 +111,12 @@ 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)
-            )
+            iam_role_arn = instance_profile.iam_role_arn
+            role_match = re.match(AWSResources.ROLE_NAME_REGEX, instance_profile.instance_profile_arn)
+            if role_match is not None:
+                instance_profile_name = role_match.group(1)
+                iam_role_arn = self._aws_resources.get_instance_profile_role_arn(instance_profile_name)
+            result_instance_profiles.append(AWSInstanceProfile(instance_profile.instance_profile_arn, iam_role_arn))
 
         return result_instance_profiles
 
@@ -230,7 +229,7 @@ def _update_sql_dac_with_instance_profile(self, iam_instance_profile: AWSInstanc
         )
 
     def get_instance_profile(self, instance_profile_name: str) -> AWSInstanceProfile | None:
-        instance_profile_arn = self._aws_resources.get_instance_profile(instance_profile_name)
+        instance_profile_arn = self._aws_resources.get_instance_profile_arn(instance_profile_name)
 
         if not instance_profile_arn:
             return None

tests/unit/aws/test_access.py

@@ -199,7 +199,7 @@ def test_create_uber_principal_existing_role_in_policy(mock_ws, mock_installatio
     )
     mock_ws.cluster_policies.get.return_value = cluster_policy
     aws = create_autospec(AWSResources)
-    aws.get_instance_profile.return_value = instance_profile_arn
+    aws.get_instance_profile_arn.return_value = instance_profile_arn
     locations = ExternalLocations(mock_ws, backend, "ucx")
     prompts = MockPrompts({"We have identified existing UCX migration role *": "yes"})
     aws_resource_permissions = AWSResourcePermissions(
@@ -228,7 +228,7 @@ def test_create_uber_principal_existing_role(mock_ws, mock_installation, backend
     )
     instance_profile_arn = "arn:aws:iam::12345:instance-profile/role1"
     aws = create_autospec(AWSResources)
-    aws.get_instance_profile.return_value = instance_profile_arn
+    aws.get_instance_profile_arn.return_value = instance_profile_arn
     locations = ExternalLocations(mock_ws, backend, "ucx")
     prompts = MockPrompts(
         {
@@ -267,7 +267,7 @@ def test_create_uber_principal_no_existing_role(mock_ws, mock_installation, back
     instance_profile_arn = "arn:aws:iam::12345:instance-profile/role1"
     aws.create_migration_role.return_value = instance_profile_arn
     aws.create_instance_profile.return_value = instance_profile_arn
-    aws.get_instance_profile.return_value = instance_profile_arn
+    aws.get_instance_profile_arn.return_value = instance_profile_arn
     locations = ExternalLocations(mock_ws, backend, "ucx")
     prompts = MockPrompts({"Do you want to create new migration role *": "yes"})
     aws_resource_permissions = AWSResourcePermissions(
@@ -460,6 +460,7 @@ def test_get_uc_compatible_roles(mock_ws, mock_installation, locations):
 
 def test_instance_profiles_empty_mapping(mock_ws, mock_installation, locations, caplog):
     aws = create_autospec(AWSResources)
+    aws.get_instance_profile_role_arn.return_value = "arn:aws:iam::12345:role/role1"
     aws_resource_permissions = AWSResourcePermissions(
         mock_installation,
         mock_ws,
@@ -488,6 +489,7 @@ def test_uc_roles_empty_mapping(mock_ws, mock_installation, locations, caplog):
 
 def test_save_instance_profile_permissions(mock_ws, mock_installation, locations):
     aws = create_autospec(AWSResources)
+    aws.get_instance_profile_role_arn.return_value = "arn:aws:iam::12345:role/role1"
     aws.get_role_policy.side_effect = [
         [
             AWSPolicyAction(
@@ -682,3 +684,60 @@ def test_save_uc_compatible_roles(mock_ws, mock_installation, locations):
             },
         ],
     )
+
+
+def test_instance_profile_lookup():
+    def instance_lookup(_):
+        ip_doc = """
+{
+    "InstanceProfile": {
+        "Path": "/",
+        "InstanceProfileName": "instance_profile_1",
+        "InstanceProfileId": "778899",
+        "Arn": "arn:aws:iam::12345678:instance-profile/instance_profile_1",
+        "CreateDate": "2024-01-01T12:00:00+00:00",
+        "Roles": [
+            {
+                "Path": "/",
+                "RoleName": "arn:aws:iam::12345678:role/role_1",
+                "RoleId": "445566",
+                "Arn": "arn:aws:iam::12345678:role/role_1",
+                "CreateDate": "2024-01-01T12:00:00+00:00"
+            }
+        ]
+    }
+}
+
+        """
+        return 0, ip_doc, ""
+
+    aws = AWSResources("profile", instance_lookup)
+    assert aws.get_instance_profile_role_arn("instance_profile_1") == "arn:aws:iam::12345678:role/role_1"
+
+
+def test_instance_profile_failed_lookup():
+    def instance_lookup(_):
+        ip_doc = ""
+        return 0, ip_doc, ""
+
+    aws = AWSResources("profile", instance_lookup)
+    assert aws.get_instance_profile_role_arn("instance_profile_1") is None
+
+
+def test_instance_profile_malformed_lookup():
+    def instance_lookup(_):
+        ip_doc = """
+{
+    "InstanceProfile": {
+        "Path": "/",
+        "InstanceProfileName": "instance_profile_1",
+        "InstanceProfileId": "778899",
+        "Arn": "arn:aws:iam::12345678:instance-profile/instance_profile_1"
+    }
+}
+
+        """
+        return 0, ip_doc, ""
+
+    aws = AWSResources("profile", instance_lookup)
+    assert aws.get_instance_profile_role_arn("instance_profile_1") is None