From bac21b3fbfe3ff5d70e3a103cfedf7d05bc32187 Mon Sep 17 00:00:00 2001
From: Alex Gaynor <alex.gaynor@gmail.com>
Date: Fri, 31 May 2024 07:19:00 -0700
Subject: [PATCH] Migrate PKCS#12 serialization with keys to Rust (#10901)

---
 .../hazmat/backends/openssl/backend.py        | 10 +----
 .../hazmat/bindings/_rust/pkcs12.pyi          |  2 +
 .../hazmat/primitives/serialization/pkcs12.py |  6 +--
 src/rust/src/pkcs12.rs                        | 37 +++++++++++++++++--
 4 files changed, 39 insertions(+), 16 deletions(-)

diff --git a/src/cryptography/hazmat/backends/openssl/backend.py b/src/cryptography/hazmat/backends/openssl/backend.py
index d00d1e4b072a..0da03896974f 100644
--- a/src/cryptography/hazmat/backends/openssl/backend.py
+++ b/src/cryptography/hazmat/backends/openssl/backend.py
@@ -398,14 +398,8 @@ def serialize_key_and_certificates_to_pkcs12(
         if name is not None:
             utils._check_bytes("name", name)
 
-        if isinstance(encryption_algorithm, serialization.NoEncryption):
-            nid_cert = -1
-            nid_key = -1
-            pkcs12_iter = 0
-            # mac_iter of 0 uses OpenSSL's default value
-            mac_iter = 0
-            mac_alg = self._ffi.NULL
-        elif isinstance(
+        assert not isinstance(encryption_algorithm, serialization.NoEncryption)
+        if isinstance(
             encryption_algorithm, serialization.BestAvailableEncryption
         ):
             # PKCS12 encryption is hopeless trash and can never be fixed.
diff --git a/src/cryptography/hazmat/bindings/_rust/pkcs12.pyi b/src/cryptography/hazmat/bindings/_rust/pkcs12.pyi
index 76dd0194c40a..dcb3fca8cf1b 100644
--- a/src/cryptography/hazmat/bindings/_rust/pkcs12.pyi
+++ b/src/cryptography/hazmat/bindings/_rust/pkcs12.pyi
@@ -8,6 +8,7 @@ from cryptography import x509
 from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes
 from cryptography.hazmat.primitives.serialization.pkcs12 import (
     PKCS12KeyAndCertificates,
+    PKCS12PrivateKeyTypes,
 )
 
 class PKCS12Certificate:
@@ -35,6 +36,7 @@ def load_pkcs12(
 ) -> PKCS12KeyAndCertificates: ...
 def serialize_key_and_certificates(
     name: bytes | None,
+    key: PKCS12PrivateKeyTypes | None,
     cert: x509.Certificate | None,
     cas: typing.Iterable[x509.Certificate | PKCS12Certificate] | None,
 ) -> bytes: ...
diff --git a/src/cryptography/hazmat/primitives/serialization/pkcs12.py b/src/cryptography/hazmat/primitives/serialization/pkcs12.py
index 17e03fbbe15c..2294b54322f9 100644
--- a/src/cryptography/hazmat/primitives/serialization/pkcs12.py
+++ b/src/cryptography/hazmat/primitives/serialization/pkcs12.py
@@ -167,10 +167,8 @@ def serialize_key_and_certificates(
     if key is None and cert is None and not cas:
         raise ValueError("You must supply at least one of key, cert, or cas")
 
-    if key is None and isinstance(
-        encryption_algorithm, serialization.NoEncryption
-    ):
-        return rust_pkcs12.serialize_key_and_certificates(name, cert, cas)
+    if isinstance(encryption_algorithm, serialization.NoEncryption):
+        return rust_pkcs12.serialize_key_and_certificates(name, key, cert, cas)
 
     from cryptography.hazmat.backends.openssl.backend import backend
 
diff --git a/src/rust/src/pkcs12.rs b/src/rust/src/pkcs12.rs
index 1b1b6ceb9f28..919c40c2ad19 100644
--- a/src/rust/src/pkcs12.rs
+++ b/src/rust/src/pkcs12.rs
@@ -246,20 +246,20 @@ enum CertificateOrPKCS12Certificate {
 }
 
 #[pyo3::prelude::pyfunction]
-#[pyo3(signature = (name, cert, cas))]
+#[pyo3(signature = (name, key, cert, cas))]
 fn serialize_key_and_certificates<'p>(
     py: pyo3::Python<'p>,
     name: Option<&[u8]>,
+    key: Option<pyo3::Bound<'_, pyo3::PyAny>>,
     cert: Option<&Certificate>,
     cas: Option<pyo3::Bound<'_, pyo3::PyAny>>,
 ) -> CryptographyResult<pyo3::Bound<'p, pyo3::types::PyBytes>> {
     let (password, mac_algorithm, mac_kdf_iter) = decode_encryption_algorithm(py)?;
 
     let mut auth_safe_contents = vec![];
-    let cert_bag_contents;
+    let (cert_bag_contents, key_bag_contents);
     let mut ca_certs = vec![];
-    assert!(cert.is_some() || cas.is_some());
-    {
+    if cert.is_some() || cas.is_some() {
         let mut cert_bags = vec![];
 
         if let Some(cert) = cert {
@@ -291,6 +291,35 @@ fn serialize_key_and_certificates<'p>(
             ))),
         });
     }
+
+    if let Some(key) = key {
+        let der = types::ENCODING_DER.get(py)?;
+        let pkcs8 = types::PRIVATE_FORMAT_PKCS8.get(py)?;
+        let no_encryption = types::NO_ENCRYPTION.get(py)?.call0()?;
+
+        let pkcs8_bytes = key
+            .call_method1(
+                pyo3::intern!(py, "private_bytes"),
+                (der, pkcs8, no_encryption),
+            )?
+            .extract::<pyo3::pybacked::PyBackedBytes>()?;
+        let pkcs8_tlv = asn1::parse_single(&pkcs8_bytes)?;
+
+        let key_bag = cryptography_x509::pkcs12::SafeBag {
+            _bag_id: asn1::DefinedByMarker::marker(),
+            bag_value: asn1::Explicit::new(cryptography_x509::pkcs12::BagValue::KeyBag(pkcs8_tlv)),
+            attributes: friendly_name_attributes(name)?,
+        };
+
+        key_bag_contents = asn1::write_single(&asn1::SequenceOfWriter::new([key_bag]))?;
+        auth_safe_contents.push(cryptography_x509::pkcs7::ContentInfo {
+            _content_type: asn1::DefinedByMarker::marker(),
+            content: cryptography_x509::pkcs7::Content::Data(Some(asn1::Explicit::new(
+                &key_bag_contents,
+            ))),
+        });
+    }
+
     let auth_safe_content = asn1::write_single(&asn1::SequenceOfWriter::new(auth_safe_contents))?;
 
     let salt = types::OS_URANDOM