Managed identity support for device-identity import and export (#344)

* Fix for new identity parameter format
* test updates
* better differentiation of storage vars
* updated to remove *all* user-identitites after storage tests until CLI core is patched

azext_iot/_help.py

@@ -256,6 +256,17 @@
     - name: Export all device identities to a configured blob container using a file path which contains the SAS uri.
       text: >
         az iot hub device-identity export -n {iothub_name} --bcu {sas_uri_filepath}
+    - name: Export all device identities to a configured blob container and include device keys. Uses system assigned identity that has
+            Storage Blob Data Contributor roles for the storage account. The blob container uri does not need the blob SAS token.
+      text: >
+        az iot hub device-identity export -n {iothub_name} --ik --bcu
+        'https://mystorageaccount.blob.core.windows.net/devices' --auth-type identity --identity [system]
+    - name: Export all device identities to a configured blob container and include device keys. Uses user assigned managed identity
+            that has Storage Blob Data Contributor roles for the storage account and contributor for the IoT hub. The blob container
+            uri does not need the blob SAS token.
+      text: >
+        az iot hub device-identity export -n {iothub_name} --ik --bcu
+        'https://mystorageaccount.blob.core.windows.net/devices' --auth-type identity --identity {managed_identity_resource_id}
 """
 
 helps[
@@ -273,6 +284,15 @@
     - name: Import all device identities from a blob using a file path which contains SAS uri.
       text: >
         az iot hub device-identity import -n {iothub_name} --ibcu {input_sas_uri_filepath} --obcu {output_sas_uri_filepath}
+    - name: Import all device identities from a blob using system assigned identity that has Storage Blob Data Contributor
+            roles for both storage accounts. The blob container uri does not need the blob SAS token.
+      text: >
+        az iot hub device-identity import -n {iothub_name} --ibcu {input_sas_uri} --obcu {output_sas_uri} --auth-type identity --identity [system]
+    - name: Import all device identities from a blob using user assigned managed identity that has Storage Blob Data Contributor
+            roles for both storage accounts and contributor for the IoT hub. The blob container uri does not need the blob SAS token.
+      text: >
+        az iot hub device-identity import -n {iothub_name} --ibcu {input_sas_uri} --obcu {output_sas_uri}
+        --auth-type identity --identity {managed_identity_resource_id}
 """
 
 helps[

azext_iot/_params.py

@@ -422,7 +422,8 @@ def load_arguments(self, _):
             help="Blob Shared Access Signature URI with write, read, and delete access to "
             "a blob container. This is used to output the status of the "
             "job and the results. Note: when using Identity-based authentication an "
-            "https:// URI is still required. Input for this argument can be inline or from a file path.",
+            "https:// URI is still required - but no SAS token is necessary. Input for this argument "
+            "can be inline or from a file path.",
         )
         context.argument(
             "include_keys",
@@ -437,6 +438,15 @@ def load_arguments(self, _):
             arg_type=get_enum_type(AuthenticationType),
             help="Authentication type for communicating with the storage container.",
         )
+        context.argument(
+            "identity",
+            options_list=["--identity"],
+            help="Managed identity type to determine if system assigned managed identity or "
+            "user assigned managed identity is used. For system assigned managed identity, use "
+            "[system]. For user assigned managed identity, provide the user assigned managed "
+            "identity resource id. This identity requires a Storage Blob Data Contributor roles for the Storage "
+            "Account.",
+        )
 
     with self.argument_context("iot hub device-identity import") as context:
         context.argument(
@@ -445,24 +455,33 @@ def load_arguments(self, _):
             help="Blob Shared Access Signature URI with read access to a blob "
             "container. This blob contains the operations to be performed on "
             "the identity registry. Note: when using Identity-based authentication "
-            "an https:// URI is still required. Input for this argument can be inline "
-            "or from a file path.",
+            "an https:// URI is still required - but no SAS token is necessary. Input for this "
+            "argument can be inline or from a file path.",
         )
         context.argument(
             "output_blob_container_uri",
             options_list=["--output-blob-container-uri", "--obcu"],
             help="Blob Shared Access Signature URI with write access "
             "to a blob container. This is used to output the status of "
             "the job and the results. Note: when using Identity-based "
-            "authentication an https:// URI is still required. Input for "
-            "this argument can be inline or from a file path.",
+            "authentication an https:// URI without the SAS token is still required. "
+            "Input for this argument can be inline or from a file path.",
         )
         context.argument(
             "storage_authentication_type",
             options_list=["--auth-type", "--storage-authentication-type"],
             arg_type=get_enum_type(AuthenticationType),
             help="Authentication type for communicating with the storage container.",
         )
+        context.argument(
+            "identity",
+            options_list=["--identity"],
+            help="Managed identity type to determine if system assigned managed identity or "
+            "user assigned managed identity is used. For system assigned managed identity, use "
+            "[system]. For user assigned managed identity, provide the user assigned managed "
+            "identity resource id. This identity requires a Storage Blob Data Contributor role for the target Storage "
+            "Account and Contributor role for the IoT Hub.",
+        )
 
     with self.argument_context("iot hub device-identity parent set") as context:
         context.argument(

azext_iot/constants.py

@@ -51,4 +51,4 @@
 CONFIG_KEY_UAMQP_EXT_VERSION = "uamqp_ext_version"
 
 # Initial Track 2 SDK version
-IOTHUB_TRACK_2_SDK_MIN_VERSION = "1.0.0"
+IOTHUB_TRACK_2_SDK_MIN_VERSION = '2.0.0'

azext_iot/operations/hub.py

@@ -16,6 +16,7 @@
     TRACING_PROPERTY,
     TRACING_ALLOWED_FOR_LOCATION,
     TRACING_ALLOWED_FOR_SKU,
+    IOTHUB_TRACK_2_SDK_MIN_VERSION,
 )
 from azext_iot.common.sas_token_auth import SasTokenAuthentication
 from azext_iot.common.shared import (
@@ -2543,10 +2544,10 @@ def iot_device_export(
     blob_container_uri,
     include_keys=False,
     storage_authentication_type=None,
+    identity=None,
     resource_group_name=None,
 ):
     from azext_iot._factory import iot_hub_service_factory
-
     client = iot_hub_service_factory(cmd.cli_ctx)
     discovery = IotHubDiscovery(cmd)
     target = discovery.get_target(
@@ -2565,11 +2566,30 @@ def iot_device_export(
             if storage_authentication_type
             else None
         )
+
         export_request = ExportDevicesRequest(
             export_blob_container_uri=blob_container_uri,
             exclude_keys=not include_keys,
             authentication_type=storage_authentication_type,
         )
+
+        user_identity = identity not in [None, '[system]']
+        if user_identity and storage_authentication_type != AuthenticationType.identityBased.name:
+            raise CLIError(
+                "Device export with user-assigned identities requires identity-based authentication [--storage-auth-type]"
+            )
+        # Track 2 CLI SDKs provide support for user-assigned identity objects
+        if ensure_iothub_sdk_min_version(IOTHUB_TRACK_2_SDK_MIN_VERSION) and user_identity:
+            from azure.mgmt.iothub.models import ManagedIdentity  # pylint: disable=no-name-in-module
+            export_request.identity = ManagedIdentity(user_assigned_identity=identity)
+
+        # if the user supplied a user-assigned identity, let them know they need a new CLI/SDK
+        elif user_identity:
+            raise CLIError(
+                "Device export with user-assigned identities requires a dependency of azure-mgmt-iothub>={}"
+                .format(IOTHUB_TRACK_2_SDK_MIN_VERSION)
+            )
+
         return client.export_devices(
             target["resourcegroup"],
             hub_name,
@@ -2594,6 +2614,7 @@ def iot_device_import(
     output_blob_container_uri,
     storage_authentication_type=None,
     resource_group_name=None,
+    identity=None,
 ):
     from azext_iot._factory import iot_hub_service_factory
 
@@ -2611,20 +2632,39 @@ def iot_device_import(
 
     if ensure_iothub_sdk_min_version("0.12.0"):
         from azure.mgmt.iothub.models import ImportDevicesRequest
+
         from azext_iot.common.shared import AuthenticationType
 
         storage_authentication_type = (
             AuthenticationType(storage_authentication_type).name
             if storage_authentication_type
             else None
         )
+
         import_request = ImportDevicesRequest(
             input_blob_container_uri=input_blob_container_uri,
             output_blob_container_uri=output_blob_container_uri,
             input_blob_name=None,
             output_blob_name=None,
             authentication_type=storage_authentication_type,
         )
+
+        user_identity = identity not in [None, '[system]']
+        if user_identity and storage_authentication_type != AuthenticationType.identityBased.name:
+            raise CLIError(
+                "Device import with user-assigned identities requires identity-based authentication [--storage-auth-type]"
+            )
+        # Track 2 CLI SDKs provide support for user-assigned identity objects
+        if ensure_iothub_sdk_min_version(IOTHUB_TRACK_2_SDK_MIN_VERSION) and user_identity:
+            from azure.mgmt.iothub.models import ManagedIdentity  # pylint: disable=no-name-in-module
+            import_request.identity = ManagedIdentity(user_assigned_identity=identity)
+        # if the user supplied a user-assigned identity, let them know they need a new CLI/SDK
+        elif user_identity:
+            raise CLIError(
+                "Device import with user-assigned identities requires a dependency of azure-mgmt-iothub>={}"
+                .format(IOTHUB_TRACK_2_SDK_MIN_VERSION)
+            )
+
         return client.import_devices(
             target["resourcegroup"],
             hub_name,
@@ -2971,7 +3011,6 @@ def _get_hub_connection_string(
                 hub.name, hub.additional_properties["resourcegroup"], policy_name
             )
         )
-
     if default_eventhub:
         cs_template_eventhub = (
             "Endpoint={};SharedAccessKeyName={};SharedAccessKey={};EntityPath={}"