From 1944197193ba457b99f61af0ebcdc8c5876120d8 Mon Sep 17 00:00:00 2001
From: Alex Gaynor <alex.gaynor@gmail.com>
Date: Mon, 3 Jun 2024 08:38:32 -0400
Subject: [PATCH] Pass encryption algorithm to Rust in PKCS#12

Extracted from #11059
---
 .../hazmat/bindings/_rust/pkcs12.pyi          |  4 +++
 .../hazmat/primitives/serialization/pkcs12.py |  4 ++-
 src/rust/src/pkcs12.rs                        | 26 ++++++++++++++-----
 3 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/src/cryptography/hazmat/bindings/_rust/pkcs12.pyi b/src/cryptography/hazmat/bindings/_rust/pkcs12.pyi
index dcb3fca8cf1b..40514c4623d5 100644
--- a/src/cryptography/hazmat/bindings/_rust/pkcs12.pyi
+++ b/src/cryptography/hazmat/bindings/_rust/pkcs12.pyi
@@ -6,6 +6,9 @@ import typing
 
 from cryptography import x509
 from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes
+from cryptography.hazmat.primitives.serialization import (
+    KeySerializationEncryption,
+)
 from cryptography.hazmat.primitives.serialization.pkcs12 import (
     PKCS12KeyAndCertificates,
     PKCS12PrivateKeyTypes,
@@ -39,4 +42,5 @@ def serialize_key_and_certificates(
     key: PKCS12PrivateKeyTypes | None,
     cert: x509.Certificate | None,
     cas: typing.Iterable[x509.Certificate | PKCS12Certificate] | None,
+    encryption_algorithm: KeySerializationEncryption,
 ) -> bytes: ...
diff --git a/src/cryptography/hazmat/primitives/serialization/pkcs12.py b/src/cryptography/hazmat/primitives/serialization/pkcs12.py
index 2294b54322f9..d1fc460d7296 100644
--- a/src/cryptography/hazmat/primitives/serialization/pkcs12.py
+++ b/src/cryptography/hazmat/primitives/serialization/pkcs12.py
@@ -168,7 +168,9 @@ def serialize_key_and_certificates(
         raise ValueError("You must supply at least one of key, cert, or cas")
 
     if isinstance(encryption_algorithm, serialization.NoEncryption):
-        return rust_pkcs12.serialize_key_and_certificates(name, key, cert, cas)
+        return rust_pkcs12.serialize_key_and_certificates(
+            name, key, cert, cas, encryption_algorithm
+        )
 
     from cryptography.hazmat.backends.openssl.backend import backend
 
diff --git a/src/rust/src/pkcs12.rs b/src/rust/src/pkcs12.rs
index 919c40c2ad19..4663b91c4e8a 100644
--- a/src/rust/src/pkcs12.rs
+++ b/src/rust/src/pkcs12.rs
@@ -230,13 +230,23 @@ fn cert_to_bag<'a>(
     })
 }
 
-fn decode_encryption_algorithm(
-    py: pyo3::Python<'_>,
-) -> CryptographyResult<(&[u8], pyo3::Bound<'_, pyo3::PyAny>, u64)> {
+fn decode_encryption_algorithm<'a>(
+    py: pyo3::Python<'a>,
+    encryption_algorithm: pyo3::Bound<'a, pyo3::PyAny>,
+) -> CryptographyResult<(
+    pyo3::pybacked::PyBackedBytes,
+    pyo3::Bound<'a, pyo3::PyAny>,
+    u64,
+)> {
     let default_hmac_alg = types::SHA256.get(py)?.call0()?;
     let default_hmac_kdf_iter = 2048;
 
-    Ok((b"", default_hmac_alg, default_hmac_kdf_iter))
+    assert!(encryption_algorithm.is_instance(&types::NO_ENCRYPTION.get(py)?)?);
+    Ok((
+        pyo3::types::PyBytes::new_bound(py, b"").extract()?,
+        default_hmac_alg,
+        default_hmac_kdf_iter,
+    ))
 }
 
 #[derive(pyo3::FromPyObject)]
@@ -246,15 +256,17 @@ enum CertificateOrPKCS12Certificate {
 }
 
 #[pyo3::prelude::pyfunction]
-#[pyo3(signature = (name, key, cert, cas))]
+#[pyo3(signature = (name, key, cert, cas, encryption_algorithm))]
 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>>,
+    encryption_algorithm: pyo3::Bound<'_, pyo3::PyAny>,
 ) -> CryptographyResult<pyo3::Bound<'p, pyo3::types::PyBytes>> {
-    let (password, mac_algorithm, mac_kdf_iter) = decode_encryption_algorithm(py)?;
+    let (password, mac_algorithm, mac_kdf_iter) =
+        decode_encryption_algorithm(py, encryption_algorithm)?;
 
     let mut auth_safe_contents = vec![];
     let (cert_bag_contents, key_bag_contents);
@@ -328,7 +340,7 @@ fn serialize_key_and_certificates<'p>(
         .extract::<pyo3::pybacked::PyBackedBytes>()?;
     let mac_algorithm_md = hashes::message_digest_from_algorithm(py, &mac_algorithm)?;
     let mac_key = pkcs12_kdf(
-        password,
+        &password,
         &salt,
         KDF_MAC_KEY_ID,
         mac_kdf_iter,