diff --git a/kmip/pie/client.py b/kmip/pie/client.py index 1bc5e50..3b881ce 100644 --- a/kmip/pie/client.py +++ b/kmip/pie/client.py @@ -102,6 +102,7 @@ class ProxyKmipClient(object): self.logger = logging.getLogger(__name__) self.attribute_factory = attributes.AttributeFactory() + self.attribute_value_factory = self.attribute_factory.value_factory self.object_factory = factory.ObjectFactory() # TODO (peter-hamilton) Consider adding validation checks for inputs. @@ -431,6 +432,46 @@ class ProxyKmipClient(object): return response_payload.unique_identifier, response_payload.attribute + @is_connected + def set_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-related fields. + Supported parameters include: + attribute_name (string): The name of the attribute being + set. Required. + attribute_value (various): The value of the attribute + being set. Required. + + Here is an example. To set an object's 'sensitive' attribute + to True, specify: + attribute_name='Sensitive' + attribute_value=True + + For a list of all supported attributes, see the + AttributeValueFactory. + + Returns: + string: The ID of the managed object the attribute was set on. + """ + a = self.attribute_value_factory.create_attribute_value_by_enum( + enums.convert_attribute_name_to_tag(kwargs.get("attribute_name")), + kwargs.get("attribute_value") + ) + request_payload = payloads.SetAttributeRequestPayload( + unique_identifier=unique_identifier, + new_attribute=cobjects.NewAttribute(attribute=a) + ) + response_payload = self.proxy.send_request_payload( + enums.Operation.SET_ATTRIBUTE, + request_payload + ) + + return response_payload.unique_identifier + @is_connected def register(self, managed_object): """ diff --git a/kmip/services/kmip_client.py b/kmip/services/kmip_client.py index 9842063..a292c99 100644 --- a/kmip/services/kmip_client.py +++ b/kmip/services/kmip_client.py @@ -341,14 +341,20 @@ class KMIPProxy(object): ) # TODO (peterhamilton) For now limit this to the new DeleteAttribute - # operation. Migrate over existing operations to use this method - # instead. + # and SetAttribute operations. Migrate over existing operations to use + # this method instead. if operation == enums.Operation.DELETE_ATTRIBUTE: if not isinstance(payload, payloads.DeleteAttributeRequestPayload): raise TypeError( "The request payload for the DeleteAttribute operation " "must be a DeleteAttributeRequestPayload object." ) + elif operation == enums.Operation.SET_ATTRIBUTE: + if not isinstance(payload, payloads.SetAttributeRequestPayload): + raise TypeError( + "The request payload for the SetAttribute operation must " + "be a SetAttributeRequestPayload object." + ) batch_item = messages.RequestBatchItem( operation=operation, @@ -388,6 +394,15 @@ class KMIPProxy(object): "Invalid response payload received for the " "DeleteAttribute operation." ) + elif batch_item.operation.value == enums.Operation.SET_ATTRIBUTE: + if not isinstance( + batch_item.response_payload, + payloads.SetAttributeRequestPayload + ): + raise exceptions.InvalidMessage( + "Invalid response payload received for the SetAttribute " + "operation." + ) return batch_item.response_payload diff --git a/kmip/tests/unit/pie/test_client.py b/kmip/tests/unit/pie/test_client.py index 6879d42..61e8628 100644 --- a/kmip/tests/unit/pie/test_client.py +++ b/kmip/tests/unit/pie/test_client.py @@ -25,6 +25,7 @@ from kmip.core import objects as obj from kmip.core.factories import attributes from kmip.core.messages import contents from kmip.core.messages import payloads +from kmip.core import primitives from kmip.core.primitives import DateTime from kmip.services.kmip_client import KMIPProxy @@ -794,6 +795,43 @@ class TestProxyKmipClient(testtools.TestCase): self.assertEqual("1", unique_identifier) self.assertIsNone(attribute) + @mock.patch( + "kmip.pie.client.KMIPProxy", + mock.MagicMock(spec_set=KMIPProxy) + ) + def test_set_attribute(self): + """ + Test that the client can set an attribute. + """ + request_payload = payloads.SetAttributeRequestPayload( + unique_identifier="1", + new_attribute=obj.NewAttribute( + attribute=primitives.Boolean( + value=True, + tag=enums.Tags.SENSITIVE + ) + ) + ) + response_payload = payloads.SetAttributeResponsePayload( + unique_identifier="1" + ) + + with ProxyKmipClient() as client: + client.proxy.send_request_payload.return_value = response_payload + + unique_identifier = client.set_attribute( + "1", + attribute_name="Sensitive", + attribute_value=True + ) + + args = ( + enums.Operation.SET_ATTRIBUTE, + request_payload + ) + 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) ) diff --git a/kmip/tests/unit/services/test_kmip_client.py b/kmip/tests/unit/services/test_kmip_client.py index 51f0847..0e99ea4 100644 --- a/kmip/tests/unit/services/test_kmip_client.py +++ b/kmip/tests/unit/services/test_kmip_client.py @@ -843,6 +843,18 @@ class TestKMIPClient(TestCase): *args ) + args = ( + OperationEnum.SET_ATTRIBUTE, + payloads.CreateRequestPayload() + ) + self.assertRaisesRegex( + TypeError, + "The request payload for the SetAttribute operation must be a " + "SetAttributeRequestPayload object.", + self.client.send_request_payload, + *args + ) + @mock.patch( "kmip.services.kmip_client.KMIPProxy._build_request_message" ) @@ -957,7 +969,6 @@ class TestKMIPClient(TestCase): attribute_index=2 ) ) - self.assertRaisesRegex( exceptions.InvalidMessage, "Invalid response payload received for the DeleteAttribute " @@ -966,6 +977,32 @@ class TestKMIPClient(TestCase): *args ) + batch_item = ResponseBatchItem( + operation=Operation(OperationEnum.SET_ATTRIBUTE), + result_status=ResultStatus(ResultStatusEnum.SUCCESS), + response_payload=response_payload + ) + send_mock.return_value = ResponseMessage(batch_items=[batch_item]) + args = ( + OperationEnum.SET_ATTRIBUTE, + payloads.SetAttributeRequestPayload( + 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 SetAttribute " + "operation.", + self.client.send_request_payload, + *args + ) + @mock.patch( "kmip.services.kmip_client.KMIPProxy._build_request_message" )