feat: Support MI auth for DPS linked-hub (#813)

cheatsheet1999 · yuelinMS · web-flow · commit a2f573029d62 · 2026-04-22T15:04:46.000-07:00
* feat: Support MI auth for DPS linked-hub

* tweak code format

* add validations to resolve comments

---------

Co-authored-by: Yuelin Zhao &lt;yuelinzhao@microsoft.com&gt;
diff --git a/azext_iot/_params.py b/azext_iot/_params.py
@@ -353,11 +353,12 @@ def load_arguments(self, _):
         context.argument(
             "hostname_type",
             options_list=["--hostname-type", "--ht"],
-            arg_type=get_enum_type(["auto", "classic", "service"]),
+            arg_type=get_enum_type(HostnameType),
             default=HostnameType.AUTO.value,
             help="Type of hostname to use in the connection string. "
-            "'auto' uses the TLS 1.3 service hostname on GWv2 hubs, classic otherwise. "
+            "'auto' uses the TLS 1.3 device hostname on GWv2 hubs, classic otherwise. "
             "'classic' always uses the default hostname. "
+            "'device' uses the TLS 1.3 device hostname (errors if not GWv2). "
             "'service' uses the TLS 1.3 service hostname (errors if not GWv2).",
         )
 
@@ -369,12 +370,13 @@ def load_arguments(self, _):
             context.argument(
                 "hostname_type",
                 options_list=["--hostname-type", "--ht"],
-                arg_type=get_enum_type(["auto", "classic", "device"]),
+                arg_type=get_enum_type(HostnameType),
                 default=HostnameType.AUTO.value,
                 help="Type of hostname to use in the connection string. "
                 "'auto' uses the TLS 1.3 device hostname on GWv2 hubs, classic otherwise. "
                 "'classic' always uses the default hostname. "
-                "'device' uses the TLS 1.3 device hostname (errors if not GWv2).",
+                "'device' uses the TLS 1.3 device hostname (errors if not GWv2). "
+                "'service' uses the TLS 1.3 service hostname (errors if not GWv2).",
             )
 
     with self.argument_context("iot hub job") as context:
diff --git a/azext_iot/constants.py b/azext_iot/constants.py
@@ -52,3 +52,4 @@
 LRO_POLL_RETRIES = 10
 
 IOTHUB_PREVIEW_API_VERSION = "2026-03-01-preview"
+IOT_HUB_DEFAULT_POLICY = "iothubowner"
diff --git a/azext_iot/core/_params.py b/azext_iot/core/_params.py
@@ -21,7 +21,7 @@
 
 
 from .custom import KeyType, SimpleAccessRights
-from .shared import IotDpsSku, IotHubSku, AccessRightsDescription
+from .shared import IotDpsSku, IotHubSku, AccessRightsDescription, IotHubAuthenticationType
 from azure.cli.command_modules.iot._validators import (validate_policy_permissions,
                                                        validate_retention_days,
                                                        validate_fileupload_notification_max_delivery_count,
@@ -134,6 +134,26 @@ def load_arguments(self, _):  # pylint: disable=too-many-statements
                    help='Location of the IoT hub.',
                    arg_group='IoT Hub Identifier',
                    deprecate_info=c.deprecate(hide=True))
+        c.argument('authentication_type',
+                   options_list=['--authentication-type', '--auth-type'],
+                   arg_type=get_enum_type(IotHubAuthenticationType),
+                   help='Authentication type for the linked IoT Hub. '
+                   "'KeyBased' uses a connection string. "
+                   "'SystemAssigned' uses the DPS system-assigned managed identity. "
+                   "'UserAssigned' uses a user-assigned managed identity.")
+        c.argument('user_assigned_identity',
+                   options_list=['--user-assigned-identity', '--uai'],
+                   help='User-assigned managed identity resource ID. '
+                   'Required when authentication type is UserAssigned.')
+        c.argument('hostname_type',
+                   options_list=['--hostname-type', '--ht'],
+                   arg_type=get_enum_type(["device", "classic"]),
+                   default="device",
+                   help="Type of IoT Hub hostname to use when linking. "
+                   "'device' uses the TLS 1.3 device hostname (hub.device.azure-devices.net). "
+                   "'classic' uses the classic hostname (hub.azure-devices.net). "
+                   "Falls back to classic if the hub does not support TLS 1.3. "
+                   "Only applies when --hub-name is provided.")
         c.argument('apply_allocation_policy',
                    help='A boolean indicating whether to apply allocation policy to the IoT hub.',
                    arg_type=get_three_state_flag())
diff --git a/azext_iot/core/custom.py b/azext_iot/core/custom.py
@@ -17,6 +17,7 @@
     BadRequestError,
     CLIInternalError,
     InvalidArgumentValueError,
+    MutuallyExclusiveArgumentError,
     RequiredArgumentMissingError,
     ResourceNotFoundError,
     UnclassifiedUserFault,
@@ -29,6 +30,8 @@
 from knack.util import CLIError
 
 from azext_iot._factory import iot_hub_service_factory, resource_service_factory
+from azext_iot.common._azure import IOT_SERVICE_CS_TEMPLATE
+from azext_iot.constants import IOT_HUB_DEFAULT_POLICY
 from azext_iot.common.certops import open_certificate
 from azext_iot.core.shared import (
     ADR_CONFIGURE_ROLES_ERROR_MSG,
@@ -41,6 +44,7 @@
     EndpointType,
     IdentityType,
     IotDpsSku,
+    IotHubAuthenticationType,
     IotHubSku,
     ManagedServiceIdentityType,
     RenewKeyType,
@@ -78,6 +82,13 @@ def _get_resource_group_from_hub(hub):
     return hub["resourcegroup"]
 
 
+def _resolve_linked_hub_hostname(hub, hostname_type="device"):
+    """Resolve IoT Hub hostname for DPS linked hub based on hostname type."""
+    if hostname_type == "device":
+        return hub["properties"].get("deviceHostName") or hub["properties"]["hostName"]
+    return hub["properties"]["hostName"]
+
+
 # CUSTOM METHODS FOR DPS
 def iot_dps_list(client, resource_group_name=None):
     if resource_group_name is None:
@@ -327,40 +338,110 @@ def iot_dps_linked_hub_create(
     connection_string=None,
     location=None,
     resource_group_name=None,
+    authentication_type=None,
+    user_assigned_identity=None,
+    hostname_type="device",
     apply_allocation_policy=None,
     allocation_weight=None,
     no_wait=False
 ):
-    if not any([connection_string, hub_name]):
-        raise RequiredArgumentMissingError("Please provide the IoT Hub name or connection string.")
-    if not connection_string:
-        # Get the connection string for the hub
-        hub_client = iot_hub_service_factory(cmd.cli_ctx)
-        connection_string = iot_hub_show_connection_string(
-            hub_client, hub_name=hub_name, resource_group_name=hub_resource_group
-        )['connectionString']
+    is_mi = authentication_type in (
+        IotHubAuthenticationType.SYSTEM_ASSIGNED.value,
+        IotHubAuthenticationType.USER_ASSIGNED.value,
+    )
 
-    if not location:
-        # Parse out hub name from connection string if needed
+    # MI based Hub Linking in DPS
+    if is_mi:
+        if connection_string:
+            raise MutuallyExclusiveArgumentError(
+                "--connection-string cannot be used with --authentication-type. "
+                "Use --hub-name instead for managed identity authentication."
+            )
         if not hub_name:
-            try:
-                hub_name = re.search(r"hostname=(.[^\;\.]+)?", connection_string, re.IGNORECASE).group(1)
-            except AttributeError:
-                raise InvalidArgumentValueError("Please provide a valid IoT Hub connection string.")
+            raise RequiredArgumentMissingError(
+                "Please provide --hub-name for managed identity authentication."
+            )
+        if authentication_type == IotHubAuthenticationType.USER_ASSIGNED.value and not user_assigned_identity:
+            raise RequiredArgumentMissingError(
+                "--user-assigned-identity is required when --authentication-type is UserAssigned."
+            )
 
         hub_client = iot_hub_service_factory(cmd.cli_ctx)
-        try:
-            location = iot_hub_get(cmd, hub_client, hub_name=hub_name, resource_group_name=hub_resource_group)["location"]
-        except CLIError:
-            raise RequiredArgumentMissingError("Please provide the IoT Hub location.")
+        hub = iot_hub_get(cmd, hub_client, hub_name=hub_name, resource_group_name=hub_resource_group)
+        host_name = _resolve_linked_hub_hostname(hub, hostname_type)
+
+        # Validate MI is enabled on DPS
+        resource_group_name = _ensure_dps_resource_group_name(client, resource_group_name, dps_name)
+        dps = iot_dps_get(client, dps_name, resource_group_name)
+        identity = dps.get("identity") or {}
+        identity_type = identity.get("type", "None") if isinstance(identity, dict) else "None"
+        if authentication_type == IotHubAuthenticationType.SYSTEM_ASSIGNED.value and "SystemAssigned" not in identity_type:
+            raise InvalidArgumentValueError(
+                f"System-assigned managed identity is not enabled on DPS '{dps_name}'. "
+                "Please enable it before linking with SystemAssigned authentication."
+            )
+        if authentication_type == IotHubAuthenticationType.USER_ASSIGNED.value and "UserAssigned" not in identity_type:
+            raise InvalidArgumentValueError(
+                f"User-assigned managed identity is not configured on DPS '{dps_name}'. "
+                "Please assign a user identity before linking with UserAssigned authentication."
+            )
 
-    resource_group_name = _ensure_dps_resource_group_name(client, resource_group_name, dps_name)
+        linked_hub_entry = {
+            "location": location or hub["location"],
+            "authenticationType": authentication_type,
+            "hostName": host_name,
+        }
+        if user_assigned_identity:
+            linked_hub_entry["selectedUserAssignedIdentityResourceId"] = user_assigned_identity
 
-    dps = iot_dps_get(client, dps_name, resource_group_name)
-    dps["properties"]["iotHubs"].append({"connectionString": connection_string,
-                                         "location": location,
-                                         "applyAllocationPolicy": apply_allocation_policy,
-                                         "allocationWeight": allocation_weight})
+    # KeyBased Hub Linking in DPS
+    else:
+        if not any([connection_string, hub_name]):
+            raise RequiredArgumentMissingError("Please provide the IoT Hub name or connection string.")
+        if not connection_string:
+            hub_client = iot_hub_service_factory(cmd.cli_ctx)
+            hub = iot_hub_get(cmd, hub_client, hub_name=hub_name, resource_group_name=hub_resource_group)
+            host_name = _resolve_linked_hub_hostname(hub, hostname_type)
+            location = location or hub["location"]
+            # Build connection string with resolved hostname
+            policies = iot_hub_policy_get(hub_client, hub_name, IOT_HUB_DEFAULT_POLICY,
+                                         _get_resource_group_from_hub(hub))
+            connection_string = IOT_SERVICE_CS_TEMPLATE.format(
+                host_name, policies["keyName"], policies["primaryKey"]
+            )
+        else:
+            if ".service.azure-devices" in connection_string.lower():
+                raise InvalidArgumentValueError(
+                    "Service hostname is not supported for DPS hub linking. "
+                    "Use a connection string with device or classic hostname."
+                )
+            if not location:
+                if not hub_name:
+                    try:
+                        hub_name = re.search(r"hostname=(.[^\;\.]+)?", connection_string, re.IGNORECASE).group(1)
+                    except AttributeError:
+                        raise InvalidArgumentValueError("Please provide a valid IoT Hub connection string.")
+
+                hub_client = iot_hub_service_factory(cmd.cli_ctx)
+                try:
+                    location = iot_hub_get(cmd, hub_client, hub_name=hub_name, resource_group_name=hub_resource_group)["location"]
+                except CLIError:
+                    raise RequiredArgumentMissingError("Please provide the IoT Hub location.")
+
+        resource_group_name = _ensure_dps_resource_group_name(client, resource_group_name, dps_name)
+        dps = iot_dps_get(client, dps_name, resource_group_name)
+
+        linked_hub_entry = {
+            "connectionString": connection_string,
+            "location": location,
+        }
+
+    if apply_allocation_policy is not None:
+        linked_hub_entry["applyAllocationPolicy"] = apply_allocation_policy
+    if allocation_weight is not None:
+        linked_hub_entry["allocationWeight"] = allocation_weight
+
+    dps["properties"]["iotHubs"].append(linked_hub_entry)
 
     if no_wait:
         return client.iot_dps_resource.begin_create_or_update(
@@ -374,8 +455,8 @@ def iot_dps_linked_hub_create(
     return iot_dps_linked_hub_list(client, dps_name, resource_group_name)
 
 
-def iot_dps_linked_hub_update(cmd, client, dps_name, linked_hub, resource_group_name=None, apply_allocation_policy=None,
-                              allocation_weight=None, no_wait=False):
+def iot_dps_linked_hub_update(cmd, client, dps_name, linked_hub, resource_group_name=None,
+                              apply_allocation_policy=None, allocation_weight=None, no_wait=False):
     if '.' not in linked_hub:
         hub_client = iot_hub_service_factory(cmd.cli_ctx)
         linked_hub = _get_iot_hub_hostname(hub_client, linked_hub)
diff --git a/azext_iot/core/shared.py b/azext_iot/core/shared.py
@@ -110,6 +110,14 @@ class IotDpsSku(Enum):
     S1 = "S1"
 
 
+class IotHubAuthenticationType(Enum):
+    """Authentication type for DPS linked IoT Hub."""
+
+    KEY_BASED = "KeyBased"
+    SYSTEM_ASSIGNED = "SystemAssigned"
+    USER_ASSIGNED = "UserAssigned"
+
+
 class AccessRightsDescription(str, Enum):
     """DPS access rights."""
 
diff --git a/azext_iot/operations/hub.py b/azext_iot/operations/hub.py
@@ -2992,9 +2992,17 @@ def _get_hub_connection_string(
         ]
 
     if hostname_type == HostnameType.AUTO.value:
-        hostname = hub["properties"].get("serviceHostName") or hub["properties"]["hostName"]
+        hostname = hub["properties"].get("deviceHostName") or hub["properties"]["hostName"]
+    elif hostname_type == HostnameType.CLASSIC.value:
+        hostname = hub["properties"]["hostName"]
     else:
-        hostname = _transform_hostname(hub["properties"]["hostName"], hostname_type)
+        key = "deviceHostName" if hostname_type == HostnameType.DEVICE.value else "serviceHostName"
+        hostname = hub["properties"].get(key)
+        if not hostname:
+            raise InvalidArgumentValueError(
+                f"The '{hostname_type}' hostname is not available for IoT Hub '{hub['name']}'. "
+                "This hostname type is only supported on GWv2 IoT Hubs."
+            )
     cs_template = "HostName={};SharedAccessKeyName={};SharedAccessKey={}"
     return [
         cs_template.format(
diff --git a/azext_iot/tests/dps/core/test_dps_linked_hub_unit.py b/azext_iot/tests/dps/core/test_dps_linked_hub_unit.py

Original file line number	Diff line number	Diff line change
`@@ -52,3 +52,4 @@`
`52`	`52`	`LRO_POLL_RETRIES = 10`
`53`	`53`
`54`	`54`	`IOTHUB_PREVIEW_API_VERSION = "2026-03-01-preview"`
	`55`	`+IOT_HUB_DEFAULT_POLICY = "iothubowner"`