Create role assignment for MSI when enable_vnet_integration is true

src/fleet/azext_fleet/_client_factory.py

@@ -4,6 +4,7 @@
 # --------------------------------------------------------------------------------------------
 
 from azure.cli.core.commands.client_factory import get_mgmt_service_client
+from azure.mgmt.msi import ManagedServiceIdentityClient
 from azure.cli.core.profiles import (
     CustomResourceType,
     ResourceType
@@ -53,3 +54,7 @@ def cf_auto_upgrade_profile_operations(cli_ctx, *_):
 def get_provider_client(cli_ctx):
     return get_mgmt_service_client(
         cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES)
+
+
+def get_msi_client(cli_ctx, subscription_id=None):
+    return get_mgmt_service_client(cli_ctx, ManagedServiceIdentityClient, subscription_id=subscription_id)

src/fleet/azext_fleet/_helpers.py

@@ -14,9 +14,11 @@
 from knack.prompting import NoTTYException, prompt_y_n
 from knack.util import CLIError
 from azure.cli.command_modules.acs._roleassignments import add_role_assignment
+from azure.mgmt.core.tools import parse_resource_id
+
 
-from azext_fleet.constants import FLEET_1P_APP_ID
 from azext_fleet._client_factory import get_provider_client
+from azext_fleet._client_factory import get_msi_client
 
 logger = get_logger(__name__)
 
@@ -154,15 +156,28 @@ def _load_kubernetes_configuration(filename):
         raise CLIError(f'Error parsing {filename} ({str(ex)})') from ex
 
 
-def assign_network_contributor_role_to_subnet(cmd, subnet_id):
+def assign_network_contributor_role_to_subnet(cmd, object_id, subnet_id):
+    if not add_role_assignment(cmd, 'Network Contributor', object_id, scope=subnet_id):
+        logger.warning("Failed to create Network Contributor role assignment on the subnet %s.\n"
+                       "This role assignment is required for the managed identity to access the subnet.\n"
+                       "Please ensure you have sufficient permissions, or ask an administrator to run:\n"
+                       "az role assignment create --assignee-principal-type ServicePrincipal --assignee-object-id %s "
+                       "--role 'Network Contributor' --scope %s",
+                       subnet_id, object_id, subnet_id)
+
+
+def get_msi_object_id(cmd, msi_resource_id):
+    parsed = parse_resource_id(msi_resource_id)
+    subscription_id = parsed['subscription']
+    resource_group_name = parsed['resource_group']
+    msi_name = parsed['resource_name']
+    msi_client = get_msi_client(cmd.cli_ctx, subscription_id=subscription_id)
+    msi = msi_client.user_assigned_identities.get(resource_name=msi_name,
+                                                  resource_group_name=resource_group_name)
+    return msi.principal_id
+
+
+def is_rp_registered(cmd):
     resource_client = get_provider_client(cmd.cli_ctx)
     provider = resource_client.providers.get("Microsoft.ContainerService")
-
-    # provider registration state being is checked to ensure that the Fleet service principal is available
-    # to create the role assignment on the subnet
-    if provider.registration_state != 'Registered':
-        raise CLIError("The Microsoft.ContainerService resource provider is not registered."
-                       "Run `az provider register -n Microsoft.ContainerService --wait`.")
-    if not add_role_assignment(cmd, 'Network Contributor', FLEET_1P_APP_ID, scope=subnet_id):
-        raise CLIError("failed to create role assignment for Fleet RP.\n"
-                       f"Do you have owner permissions on the subnet {subnet_id}?\n")
+    return provider.registration_state == 'Registered'

src/fleet/azext_fleet/_params.py

@@ -24,7 +24,8 @@
     validate_vm_size,
     validate_targets,
     validate_update_strategy_id,
-    validate_labels
+    validate_labels,
+    validate_enable_vnet_integration
 )
 
 labels_type = CLIArgumentType(
@@ -43,8 +44,8 @@ def load_arguments(self, _):
         c.argument('tags', tags_type)
         c.argument('dns_name_prefix', options_list=['--dns-name-prefix', '-p'], help='Prefix for host names that are created. If not specified, generate a host name using the managed cluster and resource group names.')
         c.argument('enable_private_cluster', action='store_true', help='Whether to create the Fleet hub as a private cluster or not.')
-        c.argument('enable_vnet_integration', action='store_true', is_preview=True, help='Whether to enable apiserver vnet integration for the Fleet hub or not.')
-        c.argument('apiserver_subnet_id', validator=validate_apiserver_subnet_id, is_preview=True, help='The subnet to be used when apiserver vnet integration is enabled.')
+        c.argument('enable_vnet_integration', validator=validate_enable_vnet_integration, action='store_true', help='Whether to enable apiserver vnet integration for the Fleet hub or not.')
+        c.argument('apiserver_subnet_id', validator=validate_apiserver_subnet_id, help='The subnet to be used when apiserver vnet integration is enabled.')
         c.argument('agent_subnet_id', validator=validate_agent_subnet_id, help='The ID of the subnet which the Fleet hub node will join on startup.')
         c.argument('enable_managed_identity', action='store_true', help='Enable system assigned managed identity (MSI) on the Fleet resource.')
         c.argument('assign_identity', validator=validate_assign_identity, help='With --enable-managed-identity, enable user assigned managed identity (MSI) on the Fleet resource by specifying the user assigned identity\'s resource Id.')

src/fleet/azext_fleet/_validators.py

@@ -60,6 +60,13 @@ def validate_assign_identity(namespace):
                 "--assign-identity is not a valid Azure resource ID.")
 
 
+def validate_enable_vnet_integration(namespace):
+    if namespace.enable_vnet_integration:
+        if not namespace.enable_managed_identity or namespace.assign_identity is None:
+            raise CLIError("--enable-vnet-integration requires user assigned managed identity to be enabled. "
+                           "Please add --enable-managed-identity and --assign-identity <identity-id> to your command.")
+
+
 def validate_targets(namespace):
     ts = namespace.targets
     if not ts:

src/fleet/azext_fleet/custom.py

@@ -11,14 +11,16 @@
 from azure.cli.core.util import sdk_no_wait, get_file_json, shell_safe_json_parse
 
 from azext_fleet._client_factory import CUSTOM_MGMT_FLEET
-from azext_fleet._helpers import print_or_merge_credentials
+from azext_fleet._helpers import is_rp_registered, print_or_merge_credentials
 from azext_fleet._helpers import assign_network_contributor_role_to_subnet
+from azext_fleet._helpers import get_msi_object_id
 from azext_fleet.constants import UPGRADE_TYPE_CONTROLPLANEONLY
 from azext_fleet.constants import UPGRADE_TYPE_FULL
 from azext_fleet.constants import UPGRADE_TYPE_NODEIMAGEONLY
 from azext_fleet.constants import UPGRADE_TYPE_ERROR_MESSAGES
 from azext_fleet.constants import SUPPORTED_GATE_STATES_FILTERS
 from azext_fleet.constants import SUPPORTED_GATE_STATES_PATCH
+from azext_fleet.constants import FLEET_1P_APP_ID
 
 
 # pylint: disable=too-many-locals
@@ -112,7 +114,17 @@ def create_fleet(cmd,
     )
 
     if enable_private_cluster:
-        assign_network_contributor_role_to_subnet(cmd, agent_subnet_id)
+        # provider registration state being is checked to ensure that the Fleet service principal is available
+        # to create the role assignment on the subnet
+        if not is_rp_registered(cmd):
+            raise CLIError("The Microsoft.ContainerService resource provider is not registered."
+                           "Run `az provider register -n Microsoft.ContainerService --wait`.")
+        assign_network_contributor_role_to_subnet(cmd, FLEET_1P_APP_ID, agent_subnet_id)
+
+    if enable_vnet_integration and assign_identity is not None:
+        object_id = get_msi_object_id(cmd, assign_identity)
+        assign_network_contributor_role_to_subnet(cmd, object_id, apiserver_subnet_id)
+        assign_network_contributor_role_to_subnet(cmd, object_id, agent_subnet_id)
 
     return sdk_no_wait(no_wait,
                        client.begin_create_or_update,