feat: add bulk key regeneration to az iot hub device-identity renew-key and az iot hub module-identity renew-key

.azure-devops/templates/run-tests-parallel.yml

@@ -63,6 +63,7 @@ steps:
 
   - ${{ if eq(parameters.runUnitTests, 'true') }}:
     - script: |
+        pip freeze
         pytest -vv ${{ parameters.path }} -k "_unit.py" --cov=azext_iot --cov-config .coveragerc --junitxml=junit/test-iotext-unit-${{ parameters.name }}.xml
       displayName: '${{ parameters.name }} unit tests'
       env:

azext_iot/_help.py

@@ -231,11 +231,16 @@
 ] = """
     type: command
     short-summary: Renew target keys of an IoT Hub device with sas authentication.
+    long-summary: Currently etags and key type `swap` are not supported for bulk key regeneration.
     examples:
       - name: Renew the primary key.
         text: az iot hub device-identity renew-key -d {device_id} -n {iothub_name} --kt primary
       - name: Swap the primary and secondary keys.
         text: az iot hub device-identity renew-key -d {device_id} -n {iothub_name} --kt swap
+      - name: Renew the secondary key for two devices and their modules.
+        text: az iot hub device-identity renew-key -d {device_id} {device_id} -n {iothub_name} --kt secondary --include-modules
+      - name: Renew the both keys for all devices within the hub.
+        text: az iot hub device-identity renew-key -d * -n {iothub_name} --kt both
 """
 
 helps[
@@ -558,11 +563,16 @@
 ] = """
     type: command
     short-summary: Renew target keys of an IoT Hub device module with sas authentication.
+    long-summary: Currently etags and key type `swap` are not supported for bulk key regeneration.
     examples:
       - name: Renew the primary key.
         text: az iot hub module-identity renew-key -m {module_name} -d {device_id} -n {iothub_name} --kt primary
       - name: Swap the primary and secondary keys.
         text: az iot hub module-identity renew-key -m {module_name} -d {device_id} -n {iothub_name} --kt swap
+      - name: Renew the secondary key for two modules.
+        text: az iot hub module-identity renew-key -m {module_name} {module_name} -d {device_id} -n {iothub_name} --kt secondary
+      - name: Renew both keys for all modules in the device.
+        text: az iot hub module-identity renew-key -m * -d {device_id} -n {iothub_name} --kt both
 """
 
 helps[

azext_iot/_params.py

@@ -467,6 +467,25 @@ def load_arguments(self, _):
             arg_type=get_enum_type(RenewKeyType),
             help="Target key type to regenerate.",
         )
+        context.argument(
+            "device_ids",
+            options_list=["--device-id", "-d"],
+            help="Space seperated list of target Device Ids. Use `*` for all devices.",
+            nargs="+",
+            action="extend"
+        )
+        context.argument(
+            "include_modules",
+            options_list=["--include-modules", "--im"],
+            help="Flag to include device modules during key regeneration.",
+            arg_type=get_three_state_flag()
+        )
+        context.argument(
+            "etag",
+            options_list=["--etag", "-e"],
+            help="Etag or entity tag corresponding to the last state of the resource. "
+            "If no etag is provided the value '*' is used. This arguement only applies to `swap`.",
+        )
 
     with self.argument_context("iot hub device-identity export") as context:
         context.argument(
@@ -614,6 +633,19 @@ def load_arguments(self, _):
             arg_type=get_enum_type(RenewKeyType),
             help="Target key type to regenerate.",
         )
+        context.argument(
+            "module_ids",
+            options_list=["--module-id", "-m"],
+            help="Space seperated list of target Module Ids. Use `*` for all modules.",
+            nargs="+",
+            action="extend"
+        )
+        context.argument(
+            "etag",
+            options_list=["--etag", "-e"],
+            help="Etag or entity tag corresponding to the last state of the resource. "
+            "If no etag is provided the value '*' is used. This arguement only applies to `swap`.",
+        )
 
     with self.argument_context("iot hub distributed-tracing update") as context:
         context.argument(

azext_iot/common/shared.py

@@ -232,6 +232,7 @@ class RenewKeyType(Enum):
     primary = KeyType.primary.value
     secondary = KeyType.secondary.value
     swap = "swap"
+    both = "both"
 
 
 class IoTHubStateType(Enum):

azext_iot/operations/hub.py

@@ -7,6 +7,7 @@
 from os.path import exists
 from knack.log import get_logger
 from enum import Enum, EnumMeta
+from tqdm import tqdm
 from azure.cli.core.azclierror import (
     ArgumentUsageError,
     CLIInternalError,
@@ -43,7 +44,6 @@
     read_file_content,
     init_monitoring,
     process_json_arg,
-    generate_key,
     generate_storage_account_sas_token,
 )
 from azext_iot._factory import SdkResolver, CloudError
@@ -498,8 +498,9 @@ def _update_device_key(target, device, auth_method, pk, sk, etag=None):
 def iot_device_key_regenerate(
     cmd,
     hub_name_or_hostname,
-    device_id,
+    device_ids,
     renew_key_type,
+    include_modules=False,
     resource_group_name=None,
     login=None,
     etag=None,
@@ -512,24 +513,55 @@ def iot_device_key_regenerate(
         login=login,
         auth_type=auth_type_dataplane,
     )
-    device = _iot_device_show(target, device_id)
-    if device["authentication"]["type"] != DeviceAuthApiType.sas.value:
-        raise ClientRequestError("Device authentication should be of type sas")
 
-    pk = device["authentication"]["symmetricKey"]["primaryKey"]
-    sk = device["authentication"]["symmetricKey"]["secondaryKey"]
-
-    if renew_key_type == RenewKeyType.primary.value:
-        pk = generate_key()
-    if renew_key_type == RenewKeyType.secondary.value:
-        sk = generate_key()
     if renew_key_type == RenewKeyType.swap.value:
+        if len(device_ids) > 1 or device_ids[0] == "*":
+            raise InvalidArgumentValueError(
+                "Currently, bulk key swap is not supported."
+            )
+        device = _iot_device_show(target, device_ids[0])
+        if device["authentication"]["type"] != DeviceAuthApiType.sas.value:
+            raise ClientRequestError("Device authentication should be of type sas")
+
+        pk = device["authentication"]["symmetricKey"]["primaryKey"]
+        sk = device["authentication"]["symmetricKey"]["secondaryKey"]
+
         temp = pk
         pk = sk
         sk = temp
+        return _update_device_key(
+            target, device, device["authentication"]["type"], pk, sk, etag
+        )
 
-    return _update_device_key(
-        target, device, device["authentication"]["type"], pk, sk, etag
+    resolver = SdkResolver(target=target)
+    service_sdk = resolver.get_sdk(SdkType.service_sdk)
+    if renew_key_type in [RenewKeyType.primary.value, RenewKeyType.secondary.value]:
+        renew_key_type += "Key"
+
+    modules = []
+    if device_ids[0] == "*":
+        devices = _iot_device_twin_list(target=target, top=None)
+        devices.extend(_iot_device_twin_list(target=target, edge_enabled=True, top=None))
+        device_ids = []
+        for device in devices:
+            if device["authenticationType"] == DeviceAuthApiType.sas.value:
+                device_ids.append(device["deviceId"])
+            # non sas devices can have sas modules...
+            if include_modules:
+                modules.extend(
+                    _iot_key_regenerate_process_modules(
+                        target=target, device_id=device["deviceId"], module_ids="*"
+                    )
+                )
+    if modules:
+        logger.info(f"Found {len(modules)} modules.")
+
+    # call friendly format
+    devices = [{"id": device_ids[i]} for i in range(len(device_ids))]
+    return _iot_key_regenerate_batch(
+        service_sdk=service_sdk,
+        renew_key_type=renew_key_type,
+        items=devices + modules,
     )
 
 
@@ -945,7 +977,7 @@ def iot_device_module_key_regenerate(
     cmd,
     hub_name_or_hostname,
     device_id,
-    module_id,
+    module_ids,
     renew_key_type,
     resource_group_name=None,
     login=None,
@@ -961,42 +993,110 @@ def iot_device_module_key_regenerate(
     )
     resolver = SdkResolver(target=target)
     service_sdk = resolver.get_sdk(SdkType.service_sdk)
-    try:
-        module = service_sdk.modules.get_identity(
-            id=device_id, mid=module_id, raw=True
-        ).response.json()
-    except CloudError as e:
-        handle_service_exception(e)
 
-    if module["authentication"]["type"] != "sas":
-        raise ClientRequestError("Module authentication should be of type sas")
+    if renew_key_type == RenewKeyType.swap.value:
+        if len(module_ids) > 1 or module_ids[0] == "*":
+            raise InvalidArgumentValueError(
+                "Currently, bulk key swap is not supported."
+            )
+        try:
+            module = service_sdk.modules.get_identity(
+                id=device_id, mid=module_ids[0], raw=True
+            ).response.json()
+        except CloudError as e:
+            handle_service_exception(e)
+        if module["authentication"]["type"] != DeviceAuthApiType.sas.value:
+            raise ClientRequestError("Module authentication should be of type sas")
 
-    pk = module["authentication"]["symmetricKey"]["primaryKey"]
-    sk = module["authentication"]["symmetricKey"]["secondaryKey"]
+        pk = module["authentication"]["symmetricKey"]["primaryKey"]
+        sk = module["authentication"]["symmetricKey"]["secondaryKey"]
 
-    if renew_key_type == RenewKeyType.primary.value:
-        pk = generate_key()
-    if renew_key_type == RenewKeyType.secondary.value:
-        sk = generate_key()
-    if renew_key_type == RenewKeyType.swap.value:
         temp = pk
-        pk = sk
-        sk = temp
+        module["authentication"]["symmetricKey"]["primaryKey"] = sk
+        module["authentication"]["symmetricKey"]["secondaryKey"] = temp
+        try:
+            return service_sdk.modules.create_or_update_identity(
+                id=device_id,
+                mid=module_ids[0],
+                module=module,
+                custom_headers={
+                    "If-Match": '"{}"'.format(etag if etag else "*")
+                },
+            )
+        except CloudError as e:
+            handle_service_exception(e)
 
-    module["authentication"]["symmetricKey"]["primaryKey"] = pk
-    module["authentication"]["symmetricKey"]["secondaryKey"] = sk
+    if renew_key_type in [RenewKeyType.primary.value, RenewKeyType.secondary.value]:
+        renew_key_type += "Key"
 
-    try:
-        headers = {}
-        headers["If-Match"] = '"{}"'.format(etag if etag else "*")
-        return service_sdk.modules.create_or_update_identity(
-            id=device_id,
-            mid=module_id,
-            module=module,
-            custom_headers=headers,
-        )
-    except CloudError as e:
-        handle_service_exception(e)
+    modules = _iot_key_regenerate_process_modules(
+        target, device_id, module_ids
+    )
+    return _iot_key_regenerate_batch(
+        service_sdk=service_sdk,
+        renew_key_type=renew_key_type,
+        items=modules,
+        device_id=device_id
+    )
+
+
+def _iot_key_regenerate_process_modules(
+    target,
+    device_id,
+    module_ids,
+):
+    if module_ids[0] == "*":
+        modules = _iot_device_module_list(target=target, device_id=device_id, top=None)
+        module_ids = []
+        for module in modules:
+            # not going to question why the call the Module object - just making things
+            # easier for unit tests
+            if not isinstance(module, dict):
+                module = module.serialize()
+            if module["authentication"]["type"] == DeviceAuthApiType.sas.value:
+                module_ids.append(module["moduleId"])
+
+    return [{"id": device_id, "moduleId": module_ids[i]} for i in range(len(module_ids))]
+
+
+def _iot_key_regenerate_batch(
+    service_sdk,
+    renew_key_type,
+    items,
+    device_id=None,
+):
+    overall_result = {
+        "policyKey": renew_key_type,
+        "errors": [],
+        "rotatedKeys": []
+    }
+    starting_msg = f"modules for device {device_id}" if device_id else "devices"
+    logger.info(f"Found {len(items)} {starting_msg}.")
+    if not items:
+        return {}
+    batches = []
+    while items:
+        if len(items) > 100:
+            batches.append(items[:100])
+            items = items[100:]
+        else:
+            batches.append(items[:])
+            items = []
+    for i in tqdm(range(len(batches)), desc="Bulk key regeneration is in progress", ascii=' #'):
+        # call
+        try:
+            result = service_sdk.service.bulk_regenerate_device_key_method(
+                policy_key=renew_key_type,
+                devices=batches[i]
+            )
+        except CloudError as e:
+            handle_service_exception(e)
+        # combine result
+        if result.errors:
+            overall_result["errors"].extend(result.errors)
+        if result.rotated_keys:
+            overall_result["rotatedKeys"].extend(result.rotated_keys)
+    return overall_result
 
 
 def iot_device_module_list(

azext_iot/sdk/iothub/service/iot_hub_gateway_service_ap_is.py

@@ -20,6 +20,7 @@
 from .operations.query_operations import QueryOperations
 from .operations.jobs_operations import JobsOperations
 from .operations.cloud_to_device_messages_operations import CloudToDeviceMessagesOperations
+from .operations.service_operations import ServiceOperations
 from .operations.modules_operations import ModulesOperations
 from .operations.digital_twin_operations import DigitalTwinOperations
 from . import models
@@ -47,6 +48,7 @@ def __init__(
         super(IotHubGatewayServiceAPIsConfiguration, self).__init__(base_url)
 
         self.add_user_agent('iothubgatewayserviceapis/{}'.format(VERSION))
+        self.add_user_agent('Azure-SDK-For-Python')
 
         self.credentials = credentials
 
@@ -58,7 +60,6 @@ class IotHubGatewayServiceAPIs(SDKClient):
     :vartype config: IotHubGatewayServiceAPIsConfiguration
 
     :ivar configuration: Configuration operations
-    :vartype configuration: service.operations.ConfigurationOperations
     :vartype configuration: azure.operations.ConfigurationOperations
     :ivar statistics: Statistics operations
     :vartype statistics: azure.operations.StatisticsOperations
@@ -72,11 +73,12 @@ class IotHubGatewayServiceAPIs(SDKClient):
     :vartype jobs: azure.operations.JobsOperations
     :ivar cloud_to_device_messages: CloudToDeviceMessages operations
     :vartype cloud_to_device_messages: azure.operations.CloudToDeviceMessagesOperations
+    :ivar service: Service operations
+    :vartype service: azure.operations.ServiceOperations
     :ivar modules: Modules operations
     :vartype modules: azure.operations.ModulesOperations
     :ivar digital_twin: DigitalTwin operations
     :vartype digital_twin: azure.operations.DigitalTwinOperations
-    :vartype digital_twin: service.operations.DigitalTwinOperations
 
     :param credentials: Credentials needed for the client to connect to Azure.
     :type credentials: :mod:`A msrestazure Credentials
@@ -91,7 +93,7 @@ def __init__(
         super(IotHubGatewayServiceAPIs, self).__init__(self.config.credentials, self.config)
 
         client_models = {k: v for k, v in models.__dict__.items() if isinstance(v, type)}
-        self.api_version = '2021-04-12'
+        self.api_version = '2024-03-31'
         self._serialize = Serializer(client_models)
         self._deserialize = Deserializer(client_models)
 
@@ -109,6 +111,8 @@ def __init__(
             self._client, self.config, self._serialize, self._deserialize)
         self.cloud_to_device_messages = CloudToDeviceMessagesOperations(
             self._client, self.config, self._serialize, self._deserialize)
+        self.service = ServiceOperations(
+            self._client, self.config, self._serialize, self._deserialize)
         self.modules = ModulesOperations(
             self._client, self.config, self._serialize, self._deserialize)
         self.digital_twin = DigitalTwinOperations(