diff --git a/pkcs11/_pkcs11.pyx b/pkcs11/_pkcs11.pyx index 4261fe4..c9f8756 100644 --- a/pkcs11/_pkcs11.pyx +++ b/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 (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 = PyMem_Malloc(value_len) - if value_ptr is NULL: - raise MemoryError() - memcpy(value_ptr, value_bytes, value_len) + if _is_template_attr(key): + template_list = AttributeList.from_template(value, attribute_mapper) + value_len = template_list.count * sizeof(CK_ATTRIBUTE) + value_ptr = 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 = PyMem_Malloc(value_len) + if value_ptr is NULL: + raise MemoryError() + memcpy(value_ptr, value_bytes, value_len) lst.data[index].pValue = value_ptr lst.data[index].ulValueLen = 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( attr.pValue, attr.ulValueLen) - ) + if _is_template_attr(attr.type): + template_lst = AttributeList.from_owned_pointer( + 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( attr.pValue, 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( + 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: diff --git a/pkcs11/constants.py b/pkcs11/constants.py index 6532d25..862854e 100644 --- a/pkcs11/constants.py +++ b/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 diff --git a/tests/test_aes.py b/tests/test_aes.py index 61c0c3a..58dda6f 100644 --- a/tests/test_aes.py +++ b/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), diff --git a/tests/test_attribute_lists.py b/tests/test_attribute_lists.py new file mode 100644 index 0000000..5ee5976 --- /dev/null +++ b/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))