Add State filtering support for the Locate operation

kmip/demos/pie/locate.py

@@ -34,6 +34,7 @@
     config = opts.config
     name = opts.name
     initial_dates = opts.initial_dates
+    state = opts.state
 
     attribute_factory = AttributeFactory()
 
@@ -71,6 +72,18 @@
                 t
             )
         )
+    if state:
+        state = getattr(enums.State, state, None)
+        if state:
+            attributes.append(
+                attribute_factory.create_attribute(
+                    enums.AttributeType.STATE,
+                    state
+                )
+            )
+        else:
+            logger.error("Invalid state provided: {}".format(opts.state))
+            sys.exit(-3)
 
     # Build the client and connect to the server
     with client.ProxyKmipClient(

kmip/demos/units/locate.py

@@ -37,6 +37,7 @@
     config = opts.config
     name = opts.name
     initial_dates = opts.initial_dates
+    state = opts.state
 
     attribute_factory = AttributeFactory()
     credential_factory = CredentialFactory()
@@ -94,6 +95,19 @@
                 t
             )
         )
+    if state:
+        state = getattr(enums.State, state, None)
+        if state:
+            attributes.append(
+                attribute_factory.create_attribute(
+                    enums.AttributeType.STATE,
+                    state
+                )
+            )
+        else:
+            logger.error("Invalid state provided: {}".format(opts.state))
+            client.close()
+            sys.exit(-1)
 
     result = client.locate(attributes=attributes, credential=credential)
     client.close()

kmip/demos/utils.py

@@ -252,6 +252,14 @@ def build_cli_parser(operation=None):
                 "'Tue Jul 23 18:39:01 2019'"
             )
         )
+        parser.add_option(
+            "--state",
+            action="store",
+            type="str",
+            default=None,
+            dest="state",
+            help="The state of the secret (e.g., PRE_ACTIVE, ACTIVE)"
+        )
     elif operation is Operation.REGISTER:
         parser.add_option(
             "-f",

kmip/services/server/engine.py

@@ -1601,6 +1601,10 @@ def _process_locate(self, payload):
         managed_objects = self._list_objects_with_access_controls(
                                 enums.Operation.LOCATE)
 
+        # TODO (ph) Do a single pass on the provided attributes and preprocess
+        # them as needed (e.g., tracking multiple 'Initial Date' values, etc).
+        # Locate needs to be able to error out if multiple singleton attributes
+        # like 'State' are provided in the same request.
         if payload.attributes:
 
             managed_objects_filtered = []
@@ -1637,6 +1641,19 @@ def _process_locate(self, payload):
                             )
                             add_object = False
                             break
+                    elif name == enums.AttributeType.STATE.value:
+                        value = value.value
+                        if value != attribute:
+                            self._logger.debug(
+                                "Failed match: "
+                                "the specified state ({}) does not match "
+                                "the object's state ({}).".format(
+                                    value.name,
+                                    attribute.name
+                                )
+                            )
+                            add_object = False
+                            break
                     elif name == enums.AttributeType.INITIAL_DATE.value:
                         initial_date["value"] = attribute
                         self._track_date_attributes(

kmip/tests/integration/services/test_integration.py

@@ -1345,6 +1345,20 @@ def test_symmetric_key_create_getattributes_locate_destroy(self):
         self.assertIn(uid_a, result.uuids)
         self.assertIn(uid_b, result.uuids)
 
+        # Test locating each key based off of its state.
+        result = self.client.locate(
+            attributes=[
+                self.attr_factory.create_attribute(
+                    enums.AttributeType.STATE,
+                    enums.State.PRE_ACTIVE
+                )
+            ]
+        )
+        self.assertEqual(ResultStatus.SUCCESS, result.result_status.value)
+        self.assertEqual(2, len(result.uuids))
+        self.assertIn(uid_a, result.uuids)
+        self.assertIn(uid_b, result.uuids)
+
         # Clean up keys
         result = self.client.destroy(uid_a)
         self.assertEqual(ResultStatus.SUCCESS, result.result_status.value)

kmip/tests/integration/services/test_proxykmipclient.py

@@ -984,6 +984,19 @@ def test_create_getattributes_locate_destroy(self):
         self.assertIn(a_id, result)
         self.assertIn(b_id, result)
 
+        # Test locating each key by its state.
+        result = self.client.locate(
+            attributes=[
+                self.attribute_factory.create_attribute(
+                    enums.AttributeType.STATE,
+                    enums.State.PRE_ACTIVE
+                )
+            ]
+        )
+        self.assertEqual(2, len(result))
+        self.assertIn(a_id, result)
+        self.assertIn(b_id, result)
+
         # Clean up the keys
         self.client.destroy(a_id)
         self.client.destroy(b_id)

kmip/tests/unit/services/server/test_engine.py

@@ -4682,6 +4682,101 @@ def test_locate_with_initial_date(self):
         self.assertEqual(len(response_payload.unique_identifiers), 1)
         self.assertIn(id_b, response_payload.unique_identifiers)
 
+    def test_locate_with_state(self):
+        """
+        Test the Locate operation when the 'State' attribute is given.
+        """
+        e = engine.KmipEngine()
+        e._data_store = self.engine
+        e._data_store_session_factory = self.session_factory
+        e._data_session = e._data_store_session_factory()
+        e._is_allowed_by_operation_policy = mock.Mock(return_value=True)
+        e._logger = mock.MagicMock()
+
+        key = (
+            b'\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00'
+        )
+
+        obj_a = pie_objects.SymmetricKey(
+            enums.CryptographicAlgorithm.AES,
+            128,
+            key,
+            name='name1'
+        )
+        obj_a.state = enums.State.PRE_ACTIVE
+        obj_b = pie_objects.SymmetricKey(
+            enums.CryptographicAlgorithm.DES,
+            128,
+            key,
+            name='name2'
+        )
+        obj_b.state = enums.State.ACTIVE
+
+        e._data_session.add(obj_a)
+        e._data_session.add(obj_b)
+        e._data_session.commit()
+        e._data_session = e._data_store_session_factory()
+
+        id_a = str(obj_a.unique_identifier)
+        id_b = str(obj_b.unique_identifier)
+
+        attribute_factory = factory.AttributeFactory()
+
+        # Locate the object with a specific state
+        attrs = [
+            attribute_factory.create_attribute(
+                enums.AttributeType.STATE,
+                enums.State.PRE_ACTIVE
+            )
+        ]
+        payload = payloads.LocateRequestPayload(attributes=attrs)
+        e._logger.reset_mock()
+        response_payload = e._process_locate(payload)
+        e._data_session.commit()
+        e._data_session = e._data_store_session_factory()
+
+        e._logger.info.assert_any_call("Processing operation: Locate")
+        e._logger.debug.assert_any_call(
+            "Locate filter matched object: {}".format(id_a)
+        )
+        e._logger.debug.assert_any_call(
+            "Failed match: "
+            "the specified state ({}) does not match "
+            "the object's state ({}).".format(
+                enums.State.PRE_ACTIVE.name,
+                enums.State.ACTIVE.name
+            )
+        )
+        self.assertEqual(len(response_payload.unique_identifiers), 1)
+        self.assertIn(id_a, response_payload.unique_identifiers)
+
+        attrs = [
+            attribute_factory.create_attribute(
+                enums.AttributeType.STATE,
+                enums.State.ACTIVE
+            )
+        ]
+        payload = payloads.LocateRequestPayload(attributes=attrs)
+        e._logger.reset_mock()
+        response_payload = e._process_locate(payload)
+        e._data_session.commit()
+        e._data_session = e._data_store_session_factory()
+
+        e._logger.info.assert_any_call("Processing operation: Locate")
+        e._logger.debug.assert_any_call(
+            "Failed match: "
+            "the specified state ({}) does not match "
+            "the object's state ({}).".format(
+                enums.State.ACTIVE.name,
+                enums.State.PRE_ACTIVE.name
+            )
+        )
+        e._logger.debug.assert_any_call(
+            "Locate filter matched object: {}".format(id_b)
+        )
+        self.assertEqual(len(response_payload.unique_identifiers), 1)
+        self.assertIn(id_b, response_payload.unique_identifiers)
+
     def test_get(self):
         """
         Test that a Get request can be processed correctly.