Add State filtering support for the Locate operation

This change updates Locate operation support in the PyKMIP server,
allowing users to filter objects based on the object's State. Unit
tests and integration tests have been added to test and verify the
correctness of this feature.

Additionally, the Locate demo scripts have also been updated to
support State filtering. Simply use the "--state" flag to specify
a State enumeration for the Locate script to filter on.

kmip/demos/pie/locate.py

@@ -34,6 +34,7 @@ if __name__ == '__main__':
     config = opts.config
     name = opts.name
     initial_dates = opts.initial_dates
+    state = opts.state
 
     attribute_factory = AttributeFactory()
 
@@ -71,6 +72,18 @@ if __name__ == '__main__':
                 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 @@ if __name__ == '__main__':
     config = opts.config
     name = opts.name
     initial_dates = opts.initial_dates
+    state = opts.state
 
     attribute_factory = AttributeFactory()
     credential_factory = CredentialFactory()
@@ -94,6 +95,19 @@ if __name__ == '__main__':
                 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 @@ class KmipEngine(object):
         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 @@ class KmipEngine(object):
                             )
                             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 @@ class TestIntegration(testtools.TestCase):
         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 @@ class TestProxyKmipClientIntegration(testtools.TestCase):
         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 @@ class TestKmipEngine(testtools.TestCase):
         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.