Support UNWRAP_TEMPLATE, DERIVE_TEMPLATE etc.

Fixes #51.

Author: MatthiasValvekens

pkcs11/_pkcs11.pyx

@@ -47,6 +47,16 @@ cdef assertRV(rv) with gil:
     raise map_rv_to_error(rv)
 
 
+cdef inline CK_BBOOL _is_template_attr(CK_ATTRIBUTE_TYPE type):
+    cdef CK_ULONG mask = 0xf
+    cdef CK_ULONG mask_result = (type ^ ATTR_TEMPLATE_ATTRIBUTE) & ~mask
+    return <CK_BBOOL> (not mask_result)
+
+
+def template_as_attribute_list(template, attribute_mapper=None):
+    return AttributeList.from_template(template, attribute_mapper or AttributeMapper())
+
+
 cdef class AttributeList:
     """
     A list of CK_ATTRIBUTE objects.
@@ -81,15 +91,22 @@ cdef class AttributeList:
         cdef bytes value_bytes
         cdef CK_CHAR * value_ptr
         cdef Py_ssize_t value_len
+        cdef AttributeList template_list
         for index, (key, value) in enumerate(template.items()):
             lst.data[index].type = key
-            value_bytes = attribute_mapper.pack_attribute(key, value)
-            value_len = len(value_bytes)
-            # copy the result into a pointer that we manage, for consistency with the other init method
-            value_ptr = <CK_CHAR *> PyMem_Malloc(value_len)
-            if value_ptr is NULL:
-                raise MemoryError()
-            memcpy(value_ptr, <CK_CHAR *> value_bytes, <size_t> value_len)
+            if _is_template_attr(key):
+                template_list = AttributeList.from_template(value, attribute_mapper)
+                value_len = <Py_ssize_t> template_list.count * sizeof(CK_ATTRIBUTE)
+                value_ptr = <CK_CHAR *> template_list.data
+                template_list.data = NULL
+            else:
+                value_bytes = attribute_mapper.pack_attribute(key, value)
+                value_len = len(value_bytes)
+                # copy the result into a pointer that we manage, for consistency with the other init method
+                value_ptr = <CK_CHAR *> PyMem_Malloc(value_len)
+                if value_ptr is NULL:
+                    raise MemoryError()
+                memcpy(value_ptr, <CK_CHAR *> value_bytes, <size_t> value_len)
             lst.data[index].pValue = <CK_CHAR *> value_ptr
             lst.data[index].ulValueLen = <CK_ULONG> value_len
 
@@ -102,10 +119,21 @@ cdef class AttributeList:
         cdef CK_ATTRIBUTE * attr
         if index < self.count:
             attr = &self.data[index]
-            return attribute_mapper.unpack_attributes(
-                attr.type,
-                PyBytes_FromStringAndSize(<char *> attr.pValue, <Py_ssize_t> attr.ulValueLen)
-            )
+            if _is_template_attr(attr.type):
+                template_lst = AttributeList.from_owned_pointer(
+                    <CK_ATTRIBUTE *> attr.pValue,
+                    attr.ulValueLen // sizeof(CK_ATTRIBUTE)
+                )
+                result = template_lst.as_dict(attribute_mapper)
+                # prevent __dealloc__ from cleaning up this memory,
+                # we just needed the .as_dict(...)
+                template_lst.data = NULL
+                return result
+            else:
+                return attribute_mapper.unpack_attributes(
+                    attr.type,
+                    PyBytes_FromStringAndSize(<char *> attr.pValue, <Py_ssize_t> attr.ulValueLen)
+                )
         else:
             raise IndexError()
 
@@ -125,12 +153,25 @@ cdef class AttributeList:
                 return self.at_index(index, attribute_mapper)
         raise KeyError(item)
 
-    def __dealloc__(self):
+    cdef _free(self):
         cdef CK_ULONG index = 0
+        cdef CK_ATTRIBUTE current
         if self.data is not NULL:
             for index in range(self.count):
-                PyMem_Free(self.data[index].pValue)
+                current = self.data[index]
+                if _is_template_attr(current.type):
+                    # ensure template lists are cleaned up
+                    AttributeList.from_owned_pointer(
+                        <CK_ATTRIBUTE *> current.pValue,
+                        current.ulValueLen // sizeof(CK_ATTRIBUTE)
+                    )._free()
+                else:
+                    PyMem_Free(current.pValue)
             PyMem_Free(self.data)
+            self.data = NULL
+
+    def __dealloc__(self):
+        self._free()
 
 
 cdef class MechanismWithParam:

pkcs11/constants.py

@@ -62,6 +62,8 @@ def __repr__(self):
 _ARRAY_ATTRIBUTE = 0x40000000
 """Attribute consists of an array of values."""
 
+ATTR_TEMPLATE_ATTRIBUTE = _ARRAY_ATTRIBUTE | 0x210
+
 
 class Attribute(IntEnum):
     """
@@ -298,9 +300,9 @@ class Attribute(IntEnum):
 
     WRAP_WITH_TRUSTED = 0x00000210
     """Key can only be wrapped with a `TRUSTED` key."""
-    WRAP_TEMPLATE = _ARRAY_ATTRIBUTE | 0x00000211
-    UNWRAP_TEMPLATE = _ARRAY_ATTRIBUTE | 0x00000212
-    DERIVE_TEMPLATE = _ARRAY_ATTRIBUTE | 0x00000213
+    WRAP_TEMPLATE = ATTR_TEMPLATE_ATTRIBUTE | 0x1
+    UNWRAP_TEMPLATE = ATTR_TEMPLATE_ATTRIBUTE | 0x2
+    DERIVE_TEMPLATE = ATTR_TEMPLATE_ATTRIBUTE | 0x3
 
     OTP_FORMAT = 0x00000220
     OTP_LENGTH = 0x00000221

tests/test_aes.py

@@ -5,7 +5,7 @@
 from parameterized import parameterized
 
 import pkcs11
-from pkcs11 import ArgumentsBad, CTRParams, GCMParams, Mechanism, PKCS11Error
+from pkcs11 import ArgumentsBad, CTRParams, GCMParams, Mechanism, PKCS11Error, TemplateInconsistent
 
 from . import FIXME, TestCase, requires
 
@@ -203,6 +203,38 @@ def test_wrap(self):
 
         self.assertEqual(key[pkcs11.Attribute.VALUE], key2[pkcs11.Attribute.VALUE])
 
+    @requires(Mechanism.AES_KEY_WRAP)
+    def test_wrap_with_unwrap_template(self):
+        wrapping_key = self.session.generate_key(
+            pkcs11.KeyType.AES,
+            128,
+            template={
+                pkcs11.Attribute.UNWRAP_TEMPLATE: {
+                    pkcs11.Attribute.EXTRACTABLE: False,
+                }
+            },
+        )
+        key = self.session.generate_key(
+            pkcs11.KeyType.AES,
+            128,
+            template={
+                pkcs11.Attribute.EXTRACTABLE: True,
+                pkcs11.Attribute.SENSITIVE: False,
+            },
+        )
+        data = wrapping_key.wrap_key(key)
+
+        with self.assertRaises(TemplateInconsistent):
+            wrapping_key.unwrap_key(
+                pkcs11.ObjectClass.SECRET_KEY,
+                pkcs11.KeyType.AES,
+                data,
+                template={
+                    # forbidden by the unwrapping template
+                    pkcs11.Attribute.EXTRACTABLE: True,
+                },
+            )
+
     @parameterized.expand(
         [
             ("POSITIVE_128_BIT", 128, 16, TestCase.assertIsNotNone),

tests/test_attribute_lists.py

@@ -0,0 +1,52 @@
+import unittest
+
+from pkcs11 import Attribute
+from pkcs11._pkcs11 import template_as_attribute_list
+from pkcs11.attributes import AttributeMapper
+
+
+class AttributeListWithTemplateTest(unittest.TestCase):
+    def test_unwrap_template_readback(self):
+        template = {
+            Attribute.SENSITIVE: True,
+            Attribute.EXTRACTABLE: False,
+            Attribute.WRAP: True,
+            Attribute.UNWRAP: True,
+            Attribute.UNWRAP_TEMPLATE: {Attribute.EXTRACTABLE: False},
+            Attribute.LABEL: "test",
+        }
+        mapper = AttributeMapper()
+        lst = template_as_attribute_list(template)
+        self.assertEqual("test", lst.get(Attribute.LABEL, mapper))
+        self.assertEqual({Attribute.EXTRACTABLE: False}, lst.get(Attribute.UNWRAP_TEMPLATE, mapper))
+        self.assertEqual(template, lst.as_dict(mapper))
+
+    def test_derive_template_readback(self):
+        template = {
+            Attribute.SENSITIVE: True,
+            Attribute.EXTRACTABLE: False,
+            Attribute.DERIVE: True,
+            Attribute.DERIVE_TEMPLATE: {Attribute.EXTRACTABLE: False},
+            Attribute.LABEL: "test",
+        }
+        mapper = AttributeMapper()
+        lst = template_as_attribute_list(template)
+        self.assertEqual("test", lst.get(Attribute.LABEL, mapper))
+        self.assertEqual({Attribute.EXTRACTABLE: False}, lst.get(Attribute.DERIVE_TEMPLATE, mapper))
+        self.assertEqual(template, lst.as_dict(mapper))
+
+    def test_nested_template(self):
+        template = {
+            Attribute.SENSITIVE: True,
+            Attribute.EXTRACTABLE: False,
+            Attribute.DERIVE: True,
+            Attribute.DERIVE_TEMPLATE: {
+                Attribute.EXTRACTABLE: False,
+                Attribute.DERIVE_TEMPLATE: {Attribute.EXTRACTABLE: False},
+            },
+            Attribute.LABEL: "test",
+        }
+        mapper = AttributeMapper()
+        lst = template_as_attribute_list(template)
+        self.assertEqual("test", lst.get(Attribute.LABEL, mapper))
+        self.assertEqual(template, lst.as_dict(mapper))