[IoT] Support IoT Hub and DPS Data Residency feature (Azure#21348)

* Update Hub and DPS SDK/API-versions

* Update to support decoding certificate byte/bytearray responses for DPS

* Support for --enforce-data-residency argument on Hub and DPS create

* Test, processor, and recording updates

* Help text update

* Recording processor fix and test re-recording to remove credscan issues

* Test updates to fix failures caused by IoT Hub control-plane update

Co-authored-by: Ryan Kelly <rykelly@microsoft.com>

src/azure-cli-core/azure/cli/core/profiles/_shared.py

@@ -230,8 +230,8 @@ def default_api_version(self):
             'subscription_diagnostic_settings': '2017-05-01-preview'
         }),
         ResourceType.MGMT_APPSERVICE: '2020-09-01',
-        ResourceType.MGMT_IOTHUB: '2021-07-01',
-        ResourceType.MGMT_IOTDPS: '2020-03-01',
+        ResourceType.MGMT_IOTHUB: '2021-07-02',
+        ResourceType.MGMT_IOTDPS: '2021-10-15',
         ResourceType.MGMT_IOTCENTRAL: '2018-09-01',
         ResourceType.MGMT_ARO: '2020-04-30',
         ResourceType.MGMT_DATABOXEDGE: '2021-02-01-preview',

src/azure-cli/azure/cli/command_modules/iot/_help.py

@@ -209,6 +209,9 @@
   - name: Create an Azure IoT Hub Device Provisioning Service with the standard pricing tier S1, in the 'eastus' region.
     text: >
         az iot dps create --name MyDps --resource-group MyResourceGroup --location eastus
+  - name: Create an Azure IoT Hub Device Provisioning Service with data residency enforced. This will disable cross-region disaster recovery.
+    text: >
+        az iot dps create --name MyDps --resource-group MyResourceGroup --edr
 """
 
 helps['iot dps delete'] = """
@@ -458,6 +461,9 @@
   - name: Create an IoT Hub with local authentication, device SAS keys, and module SAS keys all disabled
     text: >
         az iot hub create --resource-group MyResourceGroup --name MyIotHub --location westus --disable-local-auth --disable-device-sas --disable-module-sas
+  - name: Create an IoT Hub with data residency enforced. This will disable cross-region disaster recovery.
+    text: >
+        az iot hub create --resource-group MyResourceGroup --name MyIoTHub --edr
 """
 
 helps['iot hub delete'] = """

src/azure-cli/azure/cli/command_modules/iot/_params.py

@@ -72,6 +72,11 @@ def load_arguments(self, _):  # pylint: disable=too-many-statements
         c.argument('sku', arg_type=get_enum_type(IotDpsSku),
                    help='Pricing tier for the IoT Hub Device Provisioning Service.')
         c.argument('unit', help='Units in your IoT Hub Device Provisioning Service.', type=int)
+        c.argument('enable_data_residency', arg_type=get_three_state_flag(),
+                   options_list=['--enforce-data-residency', '--edr'],
+                   help='Enforce data residency for this IoT Hub Device Provisioning Service by disabling '
+                   'cross geo-pair disaster recovery. This property is immutable once set on the resource. '
+                   'Only available in select regions. Learn more at https://aka.ms/dpsdr')
 
     # To deprecate
     for subgroup in ['access-policy', 'linked-hub', 'certificate']:
@@ -374,6 +379,11 @@ def load_arguments(self, _):  # pylint: disable=too-many-statements
         c.argument('hub_name', completer=None)
         c.argument('location', get_location_type(self.cli_ctx),
                    help='Location of your IoT Hub. Default is the location of target resource group.')
+        c.argument('enable_data_residency', arg_type=get_three_state_flag(),
+                   options_list=['--enforce-data-residency', '--edr'],
+                   help='Enforce data residency for this IoT Hub by disabling cross-region disaster recovery. '
+                   'This property is immutable once set on the resource. Only available in select regions. '
+                   'Learn more at https://aka.ms/iothubdisabledr')
 
     with self.argument_context('iot hub show-connection-string') as c:
         c.argument('show_all', options_list=['--all'], help='Allow to show all shared access policies.')

src/azure-cli/azure/cli/command_modules/iot/_utils.py

@@ -57,3 +57,40 @@ def generate_key(byte_length=32):
 
     token_bytes = secrets.token_bytes(byte_length)
     return base64.b64encode(token_bytes).decode("utf8")
+
+
+def _dps_certificate_response_transform(certificate_response):
+    from azure.mgmt.iothubprovisioningservices.models import (CertificateListDescription,
+                                                              CertificateResponse,
+                                                              VerificationCodeResponse)
+    if isinstance(certificate_response, CertificateListDescription) and certificate_response.value:
+        for cert in certificate_response.value:
+            cert = _replace_certificate_bytes(cert)
+    if isinstance(certificate_response, (CertificateResponse, VerificationCodeResponse)):
+        certificate_response = _replace_certificate_bytes(certificate_response)
+    return certificate_response
+
+
+def _replace_certificate_bytes(cert_object):
+    properties = getattr(cert_object, 'properties', {})
+    body = getattr(properties, 'certificate', None)
+    if body:
+        cert_bytes = _safe_decode(body)
+        if not cert_bytes:
+            from knack.log import get_logger
+            logger = get_logger(__name__)
+            logger.warning('Certificate `%s` contains invalid unicode characters; its body was omitted from output.',
+                           cert_object.name)
+        cert_object.properties.certificate = cert_bytes
+    return cert_object
+
+
+def _safe_decode(cert_bytes):
+    if isinstance(cert_bytes, str):
+        return cert_bytes
+    if isinstance(cert_bytes, (bytearray, bytes)):
+        try:
+            return cert_bytes.decode('utf-8')
+        except UnicodeDecodeError:
+            return None
+    return None

src/azure-cli/azure/cli/command_modules/iot/commands.py

@@ -7,6 +7,7 @@
 from ._client_factory import iot_hub_service_factory
 from ._client_factory import iot_service_provisioning_factory
 from ._client_factory import iot_central_service_factory
+from ._utils import _dps_certificate_response_transform
 
 CS_DEPRECATION_INFO = 'IoT Extension (azure-iot) connection-string command (az iot hub connection-string show)'
 
@@ -80,7 +81,9 @@ def load_command_table(self, _):  # pylint: disable=too-many-statements
         g.custom_command('delete', 'iot_dps_linked_hub_delete', supports_no_wait=True)
 
     # iot dps certificate commands
-    with self.command_group('iot dps certificate', client_factory=iot_service_provisioning_factory) as g:
+    with self.command_group('iot dps certificate',
+                            client_factory=iot_service_provisioning_factory,
+                            transform=_dps_certificate_response_transform) as g:
         g.custom_command('list', 'iot_dps_certificate_list')
         g.custom_show_command('show', 'iot_dps_certificate_get')
         g.custom_command('create', 'iot_dps_certificate_create')

src/azure-cli/azure/cli/command_modules/iot/custom.py

@@ -110,11 +110,11 @@ def iot_dps_get(client, dps_name, resource_group_name=None):
     return client.iot_dps_resource.get(dps_name, resource_group_name)
 
 
-def iot_dps_create(cmd, client, dps_name, resource_group_name, location=None, sku=IotDpsSku.s1.value, unit=1, tags=None):
+def iot_dps_create(cmd, client, dps_name, resource_group_name, location=None, sku=IotDpsSku.s1.value, unit=1, tags=None, enable_data_residency=None):
     cli_ctx = cmd.cli_ctx
     _check_dps_name_availability(client.iot_dps_resource, dps_name)
     location = _ensure_location(cli_ctx, resource_group_name, location)
-    dps_property = IotDpsPropertiesDescription()
+    dps_property = IotDpsPropertiesDescription(enable_data_residency=enable_data_residency)
     dps_description = ProvisioningServiceDescription(location=location,
                                                      properties=dps_property,
                                                      sku=IotDpsSkuInfo(name=sku, capacity=unit),
@@ -410,12 +410,7 @@ def iot_dps_certificate_delete(client, dps_name, certificate_name, etag, resourc
 
 def iot_dps_certificate_gen_code(client, dps_name, certificate_name, etag, resource_group_name=None):
     resource_group_name = _ensure_dps_resource_group_name(client, resource_group_name, dps_name)
-    response = client.dps_certificate.generate_verification_code(certificate_name, etag, resource_group_name, dps_name)
-    properties = getattr(response, 'properties', {})
-    cert = getattr(properties, 'certificate', None)
-    if isinstance(cert, bytearray):
-        response.properties.certificate = response.properties.certificate.decode('utf-8')
-    return response
+    return client.dps_certificate.generate_verification_code(certificate_name, etag, resource_group_name, dps_name)
 
 
 def iot_dps_certificate_verify(client, dps_name, certificate_name, certificate_path, etag, resource_group_name=None):
@@ -498,6 +493,7 @@ def iot_hub_create(cmd, client, hub_name, resource_group_name, location=None,
                    disable_local_auth=None,
                    disable_device_sas=None,
                    disable_module_sas=None,
+                   enable_data_residency=None,
                    feedback_lock_duration=5,
                    feedback_ttl=1,
                    feedback_max_delivery_count=10,
@@ -565,6 +561,7 @@ def iot_hub_create(cmd, client, hub_name, resource_group_name, location=None,
                                   storage_endpoints=storage_endpoint_dic,
                                   cloud_to_device=cloud_to_device_properties,
                                   min_tls_version=min_tls_version,
+                                  enable_data_residency=enable_data_residency,
                                   disable_local_auth=disable_local_auth,
                                   disable_device_sas=disable_device_sas,
                                   disable_module_sas=disable_module_sas)

src/azure-cli/azure/cli/command_modules/iot/tests/latest/recording_processors.py

@@ -34,7 +34,7 @@ def _replace_string_keys(self, val):
                          .format(MOCK_KEY), val, flags=re.IGNORECASE)
         if any(['SharedAccessKey=' in val, 'sharedaccesskey=' in val]):
             # Replaces live key with `mock_key` in `SharedAccessKey=live_key` or `sharedaccesskey=live_key` string response
-            val = re.sub(r'[S|s]hared[A|a]ccess[K|k]ey=([^\*].+=)', 'SharedAccessKey={}'
+            val = re.sub(r'[S|s]hared[A|a]ccess[K|k]ey=([^\*].+=);?', 'SharedAccessKey={};'
                          .format(MOCK_KEY), val, flags=re.IGNORECASE)
         return val
 
@@ -49,6 +49,6 @@ def _replace_byte_keys(self, val):
                          .format(MOCK_KEY).encode(), val, flags=re.IGNORECASE)
         if any([b'SharedAccessKey=' in val, b'sharedaccesskey=' in val]):
             # Replaces live key with `mock_key` in `SharedAccessKey=live_key` or `sharedaccesskey=live_key` byte response
-            val = re.sub(br'[S|s]hared[A|a]ccess[K|k]ey=([^\*].+=)', 'SharedAccessKey={}'
+            val = re.sub(br'[S|s]hared[A|a]ccess[K|k]ey=([^\*].+=);?', 'SharedAccessKey={};'
                          .format(MOCK_KEY).encode(), val, flags=re.IGNORECASE)
         return val