Merge pull request #342 from OpenKMIP/feat/add-usage-mask-to-client

Add asymmetric usage mask support to the ProxyKmipClient

kmip/pie/client.py

@@ -216,7 +216,9 @@ class ProxyKmipClient(api.KmipClient):
                         length,
                         operation_policy_name=None,
                         public_name=None,
-                        private_name=None):
+                        public_usage_mask=None,
+                        private_name=None,
+                        private_usage_mask=None):
         """
         Create an asymmetric key pair on a KMIP appliance.
 
@@ -228,8 +230,14 @@ class ProxyKmipClient(api.KmipClient):
                 to use for the new key pair. Optional, defaults to None.
             public_name (string): The name to give the public key.
                                   Optional, defaults to None.
+            public_usage_mask (list): A list of CryptographicUsageMask
+                enumerations indicating how the public key should be used.
+                Optional, defaults to None.
             private_name (string): The name to give the public key.
                                    Optional, defaults to None.
+            private_usage_mask (list): A list of CryptographicUsageMask
+                enumerations indicating how the private key should be used.
+                Optional, defaults to None.
 
         Returns:
             string: The uid of the newly created public key.
@@ -257,17 +265,39 @@ class ProxyKmipClient(api.KmipClient):
 
         # Create public / private specific attributes
         public_template = None
+        names = None
         if public_name:
-            name_attr = self._build_name_attribute(name=public_name)
+            names = self._build_name_attribute(name=public_name)
+        attrs = []
+        if public_usage_mask:
+            attrs = [
+                self.attribute_factory.create_attribute(
+                    enums.AttributeType.CRYPTOGRAPHIC_USAGE_MASK,
+                    public_usage_mask
+                )
+            ]
+        if names or attrs:
             public_template = cobjects.PublicKeyTemplateAttribute(
-                names=name_attr
+                names=names,
+                attributes=attrs
             )
 
         private_template = None
+        names = None
         if private_name:
-            name_attr = self._build_name_attribute(name=private_name)
+            names = self._build_name_attribute(name=private_name)
+        attrs = []
+        if private_usage_mask:
+            attrs = [
+                self.attribute_factory.create_attribute(
+                    enums.AttributeType.CRYPTOGRAPHIC_USAGE_MASK,
+                    private_usage_mask
+                )
+            ]
+        if names or attrs:
             private_template = cobjects.PrivateKeyTemplateAttribute(
-                names=name_attr
+                names=names,
+                attributes=attrs
             )
 
         # Create the asymmetric key pair and handle the results

kmip/tests/unit/pie/test_client.py

@@ -592,6 +592,69 @@ class TestProxyKmipClient(testtools.TestCase):
                       'public_key_template_attribute': public_template}
             client.proxy.create_key_pair.assert_called_with(**kwargs)
 
+    @mock.patch('kmip.pie.client.KMIPProxy',
+                mock.MagicMock(spec_set=KMIPProxy))
+    def test_create_key_pair_with_cryptographic_usage_masks(self):
+        """
+        Test that an asymmetric key pair can be created with proper inputs,
+        specifically testing that the private / public usage masks are
+        correctly sent with the request.
+        """
+        # Create the template to test the create key pair call
+        algorithm = enums.CryptographicAlgorithm.RSA
+        length = 2048
+        algorithm_attribute = self.attribute_factory.create_attribute(
+            enums.AttributeType.CRYPTOGRAPHIC_ALGORITHM, algorithm)
+        length_attribute = self.attribute_factory.create_attribute(
+            enums.AttributeType.CRYPTOGRAPHIC_LENGTH, length)
+        mask_attribute = self.attribute_factory.create_attribute(
+            enums.AttributeType.CRYPTOGRAPHIC_USAGE_MASK,
+            [enums.CryptographicUsageMask.ENCRYPT,
+             enums.CryptographicUsageMask.DECRYPT])
+
+        private_usage_mask = self.attribute_factory.create_attribute(
+            enums.AttributeType.CRYPTOGRAPHIC_USAGE_MASK,
+            [enums.CryptographicUsageMask.SIGN]
+        )
+        public_usage_mask = self.attribute_factory.create_attribute(
+            enums.AttributeType.CRYPTOGRAPHIC_USAGE_MASK,
+            [enums.CryptographicUsageMask.VERIFY]
+        )
+
+        pair_attributes = [
+            algorithm_attribute,
+            length_attribute,
+            mask_attribute]
+
+        template = obj.CommonTemplateAttribute(attributes=pair_attributes)
+        private_template = obj.PrivateKeyTemplateAttribute(
+            attributes=[private_usage_mask])
+        public_template = obj.PublicKeyTemplateAttribute(
+            attributes=[public_usage_mask])
+
+        status = enums.ResultStatus.SUCCESS
+        result = results.CreateKeyPairResult(
+            contents.ResultStatus(status),
+            public_key_uuid=attr.PublicKeyUniqueIdentifier(
+                'aaaaaaaa-1111-2222-3333-ffffffffffff'),
+            private_key_uuid=attr.PrivateKeyUniqueIdentifier(
+                'ffffffff-3333-2222-1111-aaaaaaaaaaaa'))
+
+        with ProxyKmipClient() as client:
+            client.proxy.create_key_pair.return_value = result
+
+            _, _ = client.create_key_pair(
+                enums.CryptographicAlgorithm.RSA,
+                2048,
+                public_usage_mask=[enums.CryptographicUsageMask.VERIFY],
+                private_usage_mask=[enums.CryptographicUsageMask.SIGN]
+            )
+
+            kwargs = {'common_template_attribute': template,
+                      'private_key_template_attribute': private_template,
+                      'public_key_template_attribute': public_template}
+            client.proxy.create_key_pair.assert_called_with(**kwargs)
+
     @mock.patch('kmip.pie.client.KMIPProxy',
                 mock.MagicMock(spec_set=KMIPProxy))
     def test_create_key_pair_on_invalid_algorithm(self):