Add ML-KEM support (PKCS#11 v3.2 / FIPS 203)

pkcs11/_pkcs11.pyx

@@ -853,6 +853,28 @@ def merge_templates(default_template, *user_templates):
     }
 
 
+cdef _check_key_type(key_type):
+    if not isinstance(key_type, KeyType):
+        raise ArgumentsBad("`key_type` must be KeyType.")
+
+
+cdef _check_bit_length(value, name, required=False):
+    if value is None and not required:
+        return
+    if not isinstance(value, int):
+        raise ArgumentsBad("`%s` is the length in bits." % name)
+
+
+cdef _resolve_capabilities(key_type, capabilities):
+    if capabilities is not None:
+        return capabilities
+    try:
+        return DEFAULT_KEY_CAPABILITIES[key_type]
+    except KeyError:
+        raise ArgumentsBad("No default capabilities for this key "
+                           "type. Please specify `capabilities`.")
+
+
 cdef class Session(HasFuncList, types.Session):
     """Extend Session with implementation."""
 
@@ -971,11 +993,8 @@ cdef class Session(HasFuncList, types.Session):
     def generate_domain_parameters(self, key_type, param_length, store=False,
                                    mechanism=None, mechanism_param=None,
                                    template=None):
-        if not isinstance(key_type, KeyType):
-            raise ArgumentsBad("`key_type` must be KeyType.")
-
-        if not isinstance(param_length, int):
-            raise ArgumentsBad("`param_length` is the length in bits.")
+        _check_key_type(key_type)
+        _check_bit_length(param_length, "param_length", required=True)
 
         mech = MechanismWithParam(
             key_type, DEFAULT_PARAM_GENERATE_MECHANISMS,
@@ -999,18 +1018,9 @@ cdef class Session(HasFuncList, types.Session):
                      mechanism=None, mechanism_param=None,
                      template=None):
 
-        if not isinstance(key_type, KeyType):
-            raise ArgumentsBad("`key_type` must be KeyType.")
-
-        if key_length is not None and not isinstance(key_length, int):
-            raise ArgumentsBad("`key_length` is the length in bits.")
-
-        if capabilities is None:
-            try:
-                capabilities = DEFAULT_KEY_CAPABILITIES[key_type]
-            except KeyError:
-                raise ArgumentsBad("No default capabilities for this key "
-                                   "type. Please specify `capabilities`.")
+        _check_key_type(key_type)
+        _check_bit_length(key_length, "key_length")
+        capabilities = _resolve_capabilities(key_type, capabilities)
 
         mech = MechanismWithParam(
             key_type, DEFAULT_GENERATE_MECHANISMS,
@@ -1052,18 +1062,9 @@ cdef class Session(HasFuncList, types.Session):
                           mechanism=None, mechanism_param=None,
                           public_template=None, private_template=None):
 
-        if not isinstance(key_type, KeyType):
-            raise ArgumentsBad("`key_type` must be KeyType.")
-
-        if key_length is not None and not isinstance(key_length, int):
-            raise ArgumentsBad("`key_length` is the length in bits.")
-
-        if capabilities is None:
-            try:
-                capabilities = DEFAULT_KEY_CAPABILITIES[key_type]
-            except KeyError:
-                raise ArgumentsBad("No default capabilities for this key "
-                                   "type. Please specify `capabilities`.")
+        _check_key_type(key_type)
+        _check_bit_length(key_length, "key_length")
+        capabilities = _resolve_capabilities(key_type, capabilities)
 
         mech = MechanismWithParam(
             key_type, DEFAULT_GENERATE_MECHANISMS,
@@ -1072,6 +1073,9 @@ cdef class Session(HasFuncList, types.Session):
         public_template_ = self.attribute_mapper.public_key_template(
             id_=id, label=label, store=store, capabilities=capabilities,
         )
+        private_template_ = self.attribute_mapper.private_key_template(
+            id_=id, label=label, store=store, capabilities=capabilities,
+        )
 
         if key_type is KeyType.RSA:
             if key_length is None:
@@ -1091,11 +1095,16 @@ cdef class Session(HasFuncList, types.Session):
                     "in `public_template` (e.g. MLDSAParameterSet.ML_DSA_65)."
                 )
 
-        public_attrs = self.make_attribute_list(merge_templates(public_template_, public_template))
+        elif key_type is KeyType.ML_KEM:
+            if public_template is None or Attribute.PARAMETER_SET not in public_template:
+                raise ArgumentsBad(
+                    "ML-KEM key generation requires `Attribute.PARAMETER_SET` "
+                    "in `public_template` (e.g. MLKEMParameterSet.ML_KEM_768)."
+                )
+            public_template_[Attribute.ENCAPSULATE] = True
+            private_template_[Attribute.DECAPSULATE] = True
 
-        private_template_ = self.attribute_mapper.private_key_template(
-            id_=id, label=label, store=store, capabilities=capabilities,
-        )
+        public_attrs = self.make_attribute_list(merge_templates(public_template_, public_template))
         private_attrs = self.make_attribute_list(merge_templates(private_template_, private_template))
         return self.generate_keypair_from_attrs(public_attrs, private_attrs, mech)
 
@@ -1370,21 +1379,23 @@ cdef object make_object(Session session, CK_OBJECT_HANDLE handle) with gil:
     """
     wrapper = ObjectHandleWrapper.wrap(session, handle)
 
-    cdef CK_ATTRIBUTE_TYPE[8] attr_keys = [
+    cdef CK_ATTRIBUTE_TYPE[10] attr_keys = [
         Attribute.CLASS,
         Attribute.ENCRYPT,
         Attribute.DECRYPT,
         Attribute.SIGN,
         Attribute.VERIFY,
         Attribute.WRAP,
         Attribute.UNWRAP,
-        Attribute.DERIVE
+        Attribute.DERIVE,
+        Attribute.ENCAPSULATE,
+        Attribute.DECAPSULATE,
     ]
 
     try:
         # Determine a list of base classes to manufacture our class with
         try:
-            attributes = wrapper.get_attribute_list(&attr_keys[0], 8)
+            attributes = wrapper.get_attribute_list(&attr_keys[0], 10)
         except PKCS11Error:
             # retry fetching the flags one by one, some tokens do not implement error handling
             # on bulk fetches correctly.
@@ -1407,6 +1418,8 @@ cdef object make_object(Session session, CK_OBJECT_HANDLE handle) with gil:
                 (Attribute.WRAP, WrapMixin),
                 (Attribute.UNWRAP, UnwrapMixin),
                 (Attribute.DERIVE, DeriveMixin),
+                (Attribute.ENCAPSULATE, EncapsulateMixin),
+                (Attribute.DECAPSULATE, DecapsulateMixin),
         ):
             try:
                 if attributes.get(attribute, session.attribute_mapper):
@@ -1445,12 +1458,7 @@ class GenerateWithParametersMixin(types.DomainParameters):
                          public_template=None, private_template=None):
 
         cdef Session session = self.session
-        if capabilities is None:
-            try:
-                capabilities = DEFAULT_KEY_CAPABILITIES[self.key_type]
-            except KeyError:
-                raise ArgumentsBad("No default capabilities for this key "
-                                   "type. Please specify `capabilities`.")
+        capabilities = _resolve_capabilities(self.key_type, capabilities)
 
         mech = MechanismWithParam(self.key_type, DEFAULT_GENERATE_MECHANISMS, mechanism, mechanism_param)
 
@@ -1869,15 +1877,8 @@ class UnwrapMixin(types.UnwrapMixin):
         if not isinstance(object_class, ObjectClass):
             raise ArgumentsBad("`object_class` must be ObjectClass.")
 
-        if not isinstance(key_type, KeyType):
-            raise ArgumentsBad("`key_type` must be KeyType.")
-
-        if capabilities is None:
-            try:
-                capabilities = DEFAULT_KEY_CAPABILITIES[key_type]
-            except KeyError:
-                raise ArgumentsBad("No default capabilities for this key "
-                                   "type. Please specify `capabilities`.")
+        _check_key_type(key_type)
+        capabilities = _resolve_capabilities(key_type, capabilities)
 
         mech = MechanismWithParam(self.key_type, DEFAULT_WRAP_MECHANISMS, mechanism, mechanism_param)
 
@@ -1920,18 +1921,9 @@ class DeriveMixin(types.DeriveMixin):
                    mechanism=None, mechanism_param=None,
                    template=None):
 
-        if not isinstance(key_type, KeyType):
-            raise ArgumentsBad("`key_type` must be KeyType.")
-
-        if not isinstance(key_length, int):
-            raise ArgumentsBad("`key_length` is the length in bits.")
-
-        if capabilities is None:
-            try:
-                capabilities = DEFAULT_KEY_CAPABILITIES[key_type]
-            except KeyError:
-                raise ArgumentsBad("No default capabilities for this key "
-                                   "type. Please specify `capabilities`.")
+        _check_key_type(key_type)
+        _check_bit_length(key_length, "key_length", required=True)
+        capabilities = _resolve_capabilities(key_type, capabilities)
 
         mech = MechanismWithParam(self.key_type, DEFAULT_DERIVE_MECHANISMS, mechanism, mechanism_param)
 
@@ -1957,6 +1949,103 @@ class DeriveMixin(types.DeriveMixin):
         return make_object(session, key)
 
 
+class EncapsulateMixin(types.EncapsulateMixin):
+    """Expand EncapsulateMixin with an implementation (ML-KEM)."""
+
+    def encapsulate_key(self, key_type,
+                        key_length=None,
+                        id=None, label=None,
+                        store=False, capabilities=None,
+                        mechanism=None, mechanism_param=None,
+                        template=None):
+
+        _check_key_type(key_type)
+        capabilities = _resolve_capabilities(key_type, capabilities)
+
+        mech = MechanismWithParam(self.key_type, DEFAULT_ENCAPSULATE_MECHANISMS, mechanism, mechanism_param)
+
+        cdef Session session = self.session
+
+        if session.funclist32 == NULL:
+            raise FunctionNotSupported("C_EncapsulateKey requires a PKCS#11 v3.2+ library")
+
+        template_ = session.attribute_mapper.secret_key_template(
+            capabilities=capabilities, id_=id, label=label, store=store,
+        )
+        if key_length is not None:
+            template_[Attribute.VALUE_LEN] = key_length // 8
+        template_[Attribute.KEY_TYPE] = key_type
+        cdef AttributeList attrs = session.make_attribute_list(merge_templates(template_, template))
+        cdef CK_MECHANISM *mech_data = mech.data
+        cdef CK_OBJECT_HANDLE pub_key = self.handle
+        cdef CK_ATTRIBUTE *attr_data = attrs.data
+        cdef CK_ULONG attr_count = attrs.count
+        cdef CK_ULONG ct_len = 0
+        cdef CK_OBJECT_HANDLE ss_handle
+        cdef CK_RV retval
+
+        with nogil:
+            retval = session.funclist32.C_EncapsulateKey(
+                session.handle, mech_data, pub_key,
+                attr_data, attr_count, NULL, &ct_len, &ss_handle)
+        assertRV(retval)
+
+        cdef CK_BYTE [:] ct_buf = CK_BYTE_buffer(ct_len)
+
+        with nogil:
+            retval = session.funclist32.C_EncapsulateKey(
+                session.handle, mech_data, pub_key,
+                attr_data, attr_count, &ct_buf[0], &ct_len, &ss_handle)
+        assertRV(retval)
+
+        return (bytes(ct_buf[:ct_len]), make_object(session, ss_handle))
+
+
+class DecapsulateMixin(types.DecapsulateMixin):
+    """Expand DecapsulateMixin with an implementation (ML-KEM)."""
+
+    def decapsulate_key(self, ciphertext, key_type,
+                        key_length=None,
+                        id=None, label=None,
+                        store=False, capabilities=None,
+                        mechanism=None, mechanism_param=None,
+                        template=None):
+
+        _check_key_type(key_type)
+        capabilities = _resolve_capabilities(key_type, capabilities)
+
+        mech = MechanismWithParam(self.key_type, DEFAULT_ENCAPSULATE_MECHANISMS, mechanism, mechanism_param)
+
+        cdef Session session = self.session
+
+        if session.funclist32 == NULL:
+            raise FunctionNotSupported("C_DecapsulateKey requires a PKCS#11 v3.2+ library")
+
+        template_ = session.attribute_mapper.secret_key_template(
+            capabilities=capabilities, id_=id, label=label, store=store,
+        )
+        if key_length is not None:
+            template_[Attribute.VALUE_LEN] = key_length // 8
+        template_[Attribute.KEY_TYPE] = key_type
+        cdef AttributeList attrs = session.make_attribute_list(merge_templates(template_, template))
+        cdef CK_MECHANISM *mech_data = mech.data
+        cdef CK_OBJECT_HANDLE priv_key = self.handle
+        cdef CK_BYTE *ct_ptr = ciphertext
+        cdef CK_ULONG ct_len = <CK_ULONG> len(ciphertext)
+        cdef CK_ATTRIBUTE *attr_data = attrs.data
+        cdef CK_ULONG attr_count = attrs.count
+        cdef CK_OBJECT_HANDLE key
+        cdef CK_RV retval
+
+        with nogil:
+            retval = session.funclist32.C_DecapsulateKey(
+                session.handle, mech_data, priv_key,
+                attr_data, attr_count, ct_ptr, ct_len, &key)
+        assertRV(retval)
+
+        return make_object(session, key)
+
+
 _CLASS_MAP = {
     ObjectClass.SECRET_KEY: SecretKey,
     ObjectClass.PUBLIC_KEY: PublicKey,

pkcs11/attributes.py

@@ -49,10 +49,12 @@ def _enum(type_: type[IntEnum]) -> Handler:
     Attribute.CHECK_VALUE: handle_bytes,
     Attribute.CLASS: _enum(ObjectClass),
     Attribute.COEFFICIENT: handle_biginteger,
+    Attribute.DECAPSULATE: handle_bool,
     Attribute.DECRYPT: handle_bool,
     Attribute.DERIVE: handle_bool,
     Attribute.EC_PARAMS: handle_bytes,
     Attribute.EC_POINT: handle_bytes,
+    Attribute.ENCAPSULATE: handle_bool,
     Attribute.ENCRYPT: handle_bool,
     Attribute.END_DATE: handle_date,
     Attribute.EXPONENT_1: handle_biginteger,
@@ -214,7 +216,9 @@ def public_key_template(
     ) -> dict[Attribute, Any]:
         template = dict(self.default_public_key_template)
         _apply_capabilities(
-            template, (Attribute.ENCRYPT, Attribute.WRAP, Attribute.VERIFY), capabilities
+            template,
+            (Attribute.ENCRYPT, Attribute.WRAP, Attribute.VERIFY),
+            capabilities,
         )
         _apply_common(template, id_, label, store)
         return template

pkcs11/constants.py

@@ -349,6 +349,11 @@ class Attribute(IntEnum):
     PARAMETER_SET = 0x0000061D
     SEED = 0x00000637
 
+    ENCAPSULATE_TEMPLATE = 0x0000062A
+    DECAPSULATE_TEMPLATE = 0x0000062B
+    ENCAPSULATE = 0x00000633
+    DECAPSULATE = 0x00000634
+
     _VENDOR_DEFINED = 0x80000000
 
     def __repr__(self) -> str:
@@ -409,6 +414,11 @@ class MechanismFlag(IntFlag):
     EC_UNCOMPRESS = 0x01000000
     EC_COMPRESS = 0x02000000
 
+    ENCAPSULATE = 0x10000000
+    """Can perform ML-KEM key encapsulation."""
+    DECAPSULATE = 0x20000000
+    """Can perform ML-KEM key decapsulation."""
+
     EXTENSION = 0x80000000
 
 

pkcs11/defaults.py

@@ -22,6 +22,7 @@
     KeyType.RSA: Mechanism.RSA_PKCS_KEY_PAIR_GEN,
     KeyType.X9_42_DH: Mechanism.X9_42_DH_KEY_PAIR_GEN,
     KeyType.EC_EDWARDS: Mechanism.EC_EDWARDS_KEY_PAIR_GEN,
+    KeyType.ML_KEM: Mechanism.ML_KEM_KEY_PAIR_GEN,
     KeyType.ML_DSA: Mechanism.ML_DSA_KEY_PAIR_GEN,
     KeyType.GENERIC_SECRET: Mechanism.GENERIC_SECRET_KEY_GEN,
 }
@@ -32,6 +33,7 @@
 _ENCRYPTION: Final[MechanismFlag] = MechanismFlag.ENCRYPT | MechanismFlag.DECRYPT
 _SIGNING: Final[MechanismFlag] = MechanismFlag.SIGN | MechanismFlag.VERIFY
 _WRAPPING: Final[MechanismFlag] = MechanismFlag.WRAP | MechanismFlag.UNWRAP
+_ENCAPSULATING: Final[MechanismFlag] = MechanismFlag.ENCAPSULATE | MechanismFlag.DECAPSULATE
 
 DEFAULT_KEY_CAPABILITIES: Final[dict[KeyType, MechanismFlag | int]] = {
     KeyType.AES: _ENCRYPTION | _SIGNING | _WRAPPING,
@@ -43,6 +45,7 @@
     KeyType.RSA: _ENCRYPTION | _SIGNING | _WRAPPING,
     KeyType.GENERIC_SECRET: 0,
     KeyType.EC_EDWARDS: _SIGNING,
+    KeyType.ML_KEM: _ENCAPSULATING,
     KeyType.ML_DSA: _SIGNING,
 }
 """
@@ -83,6 +86,11 @@
 Default mechanism for wrap/unwrap.
 """
 
+DEFAULT_ENCAPSULATE_MECHANISMS: Final[dict[KeyType, Mechanism]] = {
+    KeyType.ML_KEM: Mechanism.ML_KEM,
+}
+"""Default mechanisms for ML-KEM encapsulate/decapsulate."""
+
 DEFAULT_DERIVE_MECHANISMS: Final[dict[KeyType, Mechanism]] = {
     KeyType.DH: Mechanism.DH_PKCS_DERIVE,
     KeyType.EC: Mechanism.ECDH1_DERIVE,