Add ModifyAttribute support to the client

kmip/pie/client.py

@@ -472,6 +472,44 @@ def set_attribute(self, unique_identifier=None, **kwargs):
 
         return response_payload.unique_identifier
 
+    @is_connected
+    def modify_attribute(self, unique_identifier=None, **kwargs):
+        """
+        Set an attribute on a KMIP managed object.
+
+        Args:
+            unique_identifier (string): The ID of the managed object.
+            **kwargs (various): A placeholder for attribute values used to
+                identify the attribute to modify. For KMIP 1.0 - 1.4, the
+                supported parameters are:
+                    attribute (struct): An Attribute object containing the
+                        name and index of the existing attribute and the
+                        new value for that attribute.
+                For KMIP 2.0+, the supported parameters are:
+                    current_attribute (struct): A CurrentAttribute object
+                        containing the attribute to modify. Required if the
+                        attribute is multivalued.
+                    attribute_reference (struct): A NewAttribute object
+                        containing the new attribute value. Required.
+
+        Returns:
+            string: The ID of the managed object the attribute was modified on.
+            struct: An Attribute object representing the newly modified
+                attribute. Only returned if used for KMIP 1.0 - 1.4 messages.
+        """
+        request_payload = payloads.ModifyAttributeRequestPayload(
+            unique_identifier=unique_identifier,
+            attribute=kwargs.get("attribute"),
+            current_attribute=kwargs.get("current_attribute"),
+            new_attribute=kwargs.get("new_attribute")
+        )
+        response_payload = self.proxy.send_request_payload(
+            enums.Operation.MODIFY_ATTRIBUTE,
+            request_payload
+        )
+
+        return response_payload.unique_identifier, response_payload.attribute
+
     @is_connected
     def register(self, managed_object):
         """

kmip/services/kmip_client.py

@@ -340,8 +340,8 @@ def send_request_payload(self, operation, payload, credential=None):
                 "The request payload must be a RequestPayload object."
             )
 
-        # TODO (peterhamilton) For now limit this to the new DeleteAttribute
-        # and SetAttribute operations. Migrate over existing operations to use
+        # TODO (peterhamilton) For now limit this to the new Delete/Set/Modify
+        # Attribute operations. Migrate over existing operations to use
         # this method instead.
         if operation == enums.Operation.DELETE_ATTRIBUTE:
             if not isinstance(payload, payloads.DeleteAttributeRequestPayload):
@@ -355,6 +355,12 @@ def send_request_payload(self, operation, payload, credential=None):
                     "The request payload for the SetAttribute operation must "
                     "be a SetAttributeRequestPayload object."
                 )
+        elif operation == enums.Operation.MODIFY_ATTRIBUTE:
+            if not isinstance(payload, payloads.ModifyAttributeRequestPayload):
+                raise TypeError(
+                    "The request payload for the ModifyAttribute operation "
+                    "must be a ModifyAttributeRequestPayload object."
+                )
 
         batch_item = messages.RequestBatchItem(
             operation=operation,
@@ -403,6 +409,15 @@ def send_request_payload(self, operation, payload, credential=None):
                     "Invalid response payload received for the SetAttribute "
                     "operation."
                 )
+        elif batch_item.operation.value == enums.Operation.MODIFY_ATTRIBUTE:
+            if not isinstance(
+                batch_item.response_payload,
+                payloads.ModifyAttributeRequestPayload
+            ):
+                raise exceptions.InvalidMessage(
+                    "Invalid response payload received for the "
+                    "ModifyAttribute operation."
+                )
 
         return batch_item.response_payload
 

kmip/tests/unit/pie/test_client.py

@@ -832,6 +832,48 @@ def test_set_attribute(self):
             client.proxy.send_request_payload.assert_called_with(*args)
             self.assertEqual("1", unique_identifier)
 
+    @mock.patch(
+        "kmip.pie.client.KMIPProxy",
+        mock.MagicMock(spec_set=KMIPProxy)
+    )
+    def test_modify_attribute(self):
+        """
+        Test that the client can modify an attribute.
+        """
+        request_payload = payloads.ModifyAttributeRequestPayload(
+            unique_identifier="1",
+            new_attribute=obj.NewAttribute(
+                attribute=primitives.Boolean(
+                    value=True,
+                    tag=enums.Tags.SENSITIVE
+                )
+            )
+        )
+        response_payload = payloads.ModifyAttributeResponsePayload(
+            unique_identifier="1"
+        )
+
+        with ProxyKmipClient() as client:
+            client.proxy.send_request_payload.return_value = response_payload
+
+            unique_identifier, attribute = client.modify_attribute(
+                unique_identifier="1",
+                new_attribute=obj.NewAttribute(
+                    attribute=primitives.Boolean(
+                        value=True,
+                        tag=enums.Tags.SENSITIVE
+                    )
+                )
+            )
+
+            args = (
+                enums.Operation.MODIFY_ATTRIBUTE,
+                request_payload
+            )
+            client.proxy.send_request_payload.assert_called_with(*args)
+            self.assertEqual("1", unique_identifier)
+            self.assertIsNone(attribute)
+
     @mock.patch(
         'kmip.pie.client.KMIPProxy', mock.MagicMock(spec_set=KMIPProxy)
     )

kmip/tests/unit/services/test_kmip_client.py

@@ -855,6 +855,18 @@ def test_send_request_payload_mismatch_operation_payload(self):
             *args
         )
 
+        args = (
+            OperationEnum.MODIFY_ATTRIBUTE,
+            payloads.CreateRequestPayload()
+        )
+        self.assertRaisesRegex(
+            TypeError,
+            "The request payload for the ModifyAttribute operation must be a "
+            "ModifyAttributeRequestPayload object.",
+            self.client.send_request_payload,
+            *args
+        )
+
     @mock.patch(
         "kmip.services.kmip_client.KMIPProxy._build_request_message"
     )
@@ -977,6 +989,7 @@ def test_send_request_payload_mismatch_response_payload(
             *args
         )
 
+        # Test SetAttribute
         batch_item = ResponseBatchItem(
             operation=Operation(OperationEnum.SET_ATTRIBUTE),
             result_status=ResultStatus(ResultStatusEnum.SUCCESS),
@@ -1003,6 +1016,33 @@ def test_send_request_payload_mismatch_response_payload(
             *args
         )
 
+        # Test ModifyAttribute
+        batch_item = ResponseBatchItem(
+            operation=Operation(OperationEnum.MODIFY_ATTRIBUTE),
+            result_status=ResultStatus(ResultStatusEnum.SUCCESS),
+            response_payload=response_payload
+        )
+        send_mock.return_value = ResponseMessage(batch_items=[batch_item])
+        args = (
+            OperationEnum.MODIFY_ATTRIBUTE,
+            payloads.ModifyAttributeRequestPayload(
+                unique_identifier="1",
+                new_attribute=objects.NewAttribute(
+                    attribute=primitives.Boolean(
+                        value=True,
+                        tag=enums.Tags.SENSITIVE
+                    )
+                )
+            )
+        )
+        self.assertRaisesRegex(
+            exceptions.InvalidMessage,
+            "Invalid response payload received for the ModifyAttribute "
+            "operation.",
+            self.client.send_request_payload,
+            *args
+        )
+
     @mock.patch(
         "kmip.services.kmip_client.KMIPProxy._build_request_message"
     )