Implement convenience arguments for device update (#250)

HISTORY.rst

@@ -3,6 +3,11 @@
 Release History
 ===============
 
+**IoT Hub updates**
+
+* Add convenience arguments for device update.
+
+
 0.10.0
 +++++++++++++++
 

azext_iot/_help.py

@@ -200,9 +200,15 @@
       text: >
         az iot hub device-identity update -d {device_id} -n {iothub_name}
         --set capabilities.iotEdge=true
+    - name: Turn on edge capabilities for device using convenience argument.
+      text: >
+        az iot hub device-identity update -d {device_id} -n {iothub_name} --ee
     - name: Disable device status
       text: >
         az iot hub device-identity update -d {device_id} -n {iothub_name} --set status=disabled
+    - name: Disable device status using convenience argument.
+      text: >
+        az iot hub device-identity update -d {device_id} -n {iothub_name} --status disabled
     - name: In one command
       text: >
         az iot hub device-identity update -d {device_id} -n {iothub_name}

azext_iot/_params.py

@@ -370,6 +370,18 @@ def load_arguments(self, _):
             help="Description for device status.",
         )
 
+    with self.argument_context('iot hub device-identity update') as context:
+        context.argument(
+            "primary_key",
+            options_list=["--primary-key", "--pk"],
+            help="The primary symmetric shared access key stored in base64 format.",
+        )
+        context.argument(
+            "secondary_key",
+            options_list=["--secondary-key", "--sk"],
+            help="The secondary symmetric shared access key stored in base64 format.",
+        )
+
     with self.argument_context("iot hub device-identity create") as context:
         context.argument(
             "force",

azext_iot/commands.py

@@ -34,7 +34,11 @@ def load_command_table(self, _):
         cmd_group.command("list", "iot_device_list")
         cmd_group.command("delete", "iot_device_delete")
         cmd_group.generic_update_command(
-            "update", getter_name="iot_device_show", setter_name="iot_device_update"
+            "update",
+            getter_name="iot_device_show",
+            custom_func_type=iothub_ops,
+            setter_name="iot_device_update",
+            custom_func_name="update_iot_device_custom"
         )
 
         cmd_group.command(

azext_iot/operations/hub.py

@@ -246,6 +246,38 @@ def _create_self_signed_cert(subject, valid_days, output_path=None):
     return create_self_signed_certificate(subject, valid_days, output_path)
 
 
+def update_iot_device_custom(instance, edge_enabled=None, status=None, status_reason=None,
+                             auth_method=None, primary_thumbprint=None, secondary_thumbprint=None,
+                             primary_key=None, secondary_key=None):
+    if edge_enabled is not None:
+        instance['capabilities']['iotEdge'] = edge_enabled
+    if status is not None:
+        instance['status'] = status
+    if status_reason is not None:
+        instance['statusReason'] = status_reason
+    if auth_method is not None:
+        if auth_method == DeviceAuthType.shared_private_key.name:
+            auth = 'sas'
+            if (primary_key and not secondary_key) or (not primary_key and secondary_key):
+                raise CLIError("primary + secondary Key required with sas auth")
+            instance['authentication']['symmetricKey']['primaryKey'] = primary_key
+            instance['authentication']['symmetricKey']['secondaryKey'] = secondary_key
+        elif auth_method == DeviceAuthType.x509_thumbprint.name:
+            auth = 'selfSigned'
+            if not any([primary_thumbprint, secondary_thumbprint]):
+                raise CLIError("primary or secondary Thumbprint required with selfSigned auth")
+            if primary_thumbprint:
+                instance['authentication']['x509Thumbprint']['primaryThumbprint'] = primary_thumbprint
+            if secondary_thumbprint:
+                instance['authentication']['x509Thumbprint']['secondaryThumbprint'] = secondary_thumbprint
+        elif auth_method == DeviceAuthType.x509_ca.name:
+            auth = 'certificateAuthority'
+        else:
+            raise ValueError('Authorization method {} invalid.'.format(auth_method))
+        instance['authentication']['type'] = auth
+    return instance
+
+
 def iot_device_update(
     cmd, device_id, parameters, hub_name=None, resource_group_name=None, login=None
 ):
@@ -257,7 +289,15 @@ def iot_device_update(
     service_sdk = resolver.get_sdk(SdkType.service_sdk)
 
     try:
-        updated_device = _handle_device_update_params(parameters)
+        auth, pk, sk = _parse_auth(parameters)
+        updated_device = _assemble_device(
+            parameters['deviceId'],
+            auth, parameters['capabilities']['iotEdge'],
+            pk,
+            sk,
+            parameters['status'].lower(),
+            parameters.get('statusReason')
+        )
         etag = parameters.get("etag", None)
         if etag:
             headers = {}
@@ -272,28 +312,6 @@ def iot_device_update(
         raise CLIError(err)
 
 
-def _handle_device_update_params(parameters):
-    status = parameters["status"].lower()
-    possible_status = ["enabled", "disabled"]
-    if status not in possible_status:
-        raise CLIError("status must be one of {}".format(possible_status))
-
-    edge = parameters["capabilities"].get("iotEdge")
-    if not isinstance(edge, bool):
-        raise CLIError("capabilities.iotEdge is of type bool")
-
-    auth, pk, sk = _parse_auth(parameters)
-    return _assemble_device(
-        parameters["deviceId"],
-        auth,
-        edge,
-        pk,
-        sk,
-        status,
-        parameters.get("statusReason"),
-    )
-
-
 def iot_device_delete(
     cmd, device_id, hub_name=None, resource_group_name=None, login=None
 ):

azext_iot/tests/test_iot_ext_int.py

@@ -444,8 +444,55 @@ def test_hub_devices(self):
         )
 
         self.cmd(
-            '''iot hub device-identity update -d {} -n {} -g {} --set authentication.symmetricKey.primaryKey=""
-                    authentication.symmetricKey.secondaryKey=""'''.format(
+            "iot hub device-identity update -d {} -n {} -g {} --ee {} --auth-method {}"
+            .format(device_ids[0], LIVE_HUB, LIVE_RG, False, 'x509_ca'),
+            checks=[
+                self.check("deviceId", device_ids[0]),
+                self.check("status", "enabled"),
+                self.check("capabilities.iotEdge", False),
+                self.check("authentication.symmetricKey.primaryKey", None),
+                self.check("authentication.symmetricKey.secondaryKey", None),
+                self.check("authentication.x509Thumbprint.primaryThumbprint", None),
+                self.check("authentication.x509Thumbprint.secondaryThumbprint", None),
+                self.check("authentication.type", 'certificateAuthority')
+            ]
+        )
+
+        self.cmd(
+            "iot hub device-identity update -d {} -n {} -g {} --status-reason {}"
+            .format(device_ids[0], LIVE_HUB, LIVE_RG, 'TestStatusReason'),
+            checks=[
+                self.check("deviceId", device_ids[0]),
+                self.check("statusReason", 'TestStatusReason'),
+            ]
+        )
+
+        self.cmd(
+            "iot hub device-identity update -d {} -n {} -g {} --ee {} --status {}"
+            " --status-reason {} --auth-method {} --ptp {} --stp {}"
+            .format(device_ids[0], LIVE_HUB, LIVE_RG, False, 'enabled',
+                    'StatusReasonUpdated', 'x509_thumbprint', PRIMARY_THUMBPRINT, SECONDARY_THUMBPRINT),
+            checks=[
+                self.check("deviceId", device_ids[0]),
+                self.check("status", "enabled"),
+                self.check("capabilities.iotEdge", False),
+                self.check("statusReason", 'StatusReasonUpdated'),
+                self.check("authentication.x509Thumbprint.primaryThumbprint", PRIMARY_THUMBPRINT),
+                self.check("authentication.x509Thumbprint.secondaryThumbprint", SECONDARY_THUMBPRINT),
+            ]
+        )
+
+        self.cmd("iot hub device-identity update -d {} -n {} -g {} --auth-method {}"
+                 .format(device_ids[0], LIVE_HUB, LIVE_RG, 'x509_thumbprint'),
+                 expect_failure=True)
+
+        self.cmd("iot hub device-identity update -d {} -n {} -g {} --auth-method {} --pk {}"
+                 .format(device_ids[0], LIVE_HUB, LIVE_RG, 'shared_private_key', '123'),
+                 expect_failure=True)
+
+        self.cmd(
+            '''iot hub device-identity update -d {} -n {} -g {} --primary-key=""
+                    --secondary-key=""'''.format(
                 edge_device_ids[1], LIVE_HUB, LIVE_RG
             ),
             checks=[