From 6f09c973d9cad0dd897c34045fbf88168ea29717 Mon Sep 17 00:00:00 2001
From: Paul Kehrer <paul.l.kehrer@gmail.com>
Date: Wed, 10 Jul 2024 12:46:55 -0500
Subject: [PATCH] migrate dh, dsa, and ec to declarative (#11247)

---
 src/rust/src/backend/dh.rs  | 25 ++++++++-----------------
 src/rust/src/backend/dsa.rs | 22 +++++++---------------
 src/rust/src/backend/ec.rs  | 24 ++++++++----------------
 src/rust/src/backend/mod.rs |  3 ---
 src/rust/src/lib.rs         |  6 ++++++
 5 files changed, 29 insertions(+), 51 deletions(-)

diff --git a/src/rust/src/backend/dh.rs b/src/rust/src/backend/dh.rs
index e615d623ffa3..883277e35017 100644
--- a/src/rust/src/backend/dh.rs
+++ b/src/rust/src/backend/dh.rs
@@ -8,7 +8,7 @@ use crate::asn1::encode_der_data;
 use crate::backend::utils;
 use crate::error::{CryptographyError, CryptographyResult};
 use crate::{types, x509};
-use pyo3::types::{PyAnyMethods, PyModuleMethods};
+use pyo3::types::PyAnyMethods;
 
 const MIN_MODULUS_SIZE: u32 = 512;
 
@@ -554,20 +554,11 @@ impl DHParameterNumbers {
     }
 }
 
-pub(crate) fn create_module(
-    py: pyo3::Python<'_>,
-) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::types::PyModule>> {
-    let m = pyo3::types::PyModule::new_bound(py, "dh")?;
-    m.add_function(pyo3::wrap_pyfunction_bound!(generate_parameters, &m)?)?;
-    m.add_function(pyo3::wrap_pyfunction_bound!(from_der_parameters, &m)?)?;
-    m.add_function(pyo3::wrap_pyfunction_bound!(from_pem_parameters, &m)?)?;
-
-    m.add_class::<DHPrivateKey>()?;
-    m.add_class::<DHPublicKey>()?;
-    m.add_class::<DHParameters>()?;
-    m.add_class::<DHPrivateNumbers>()?;
-    m.add_class::<DHPublicNumbers>()?;
-    m.add_class::<DHParameterNumbers>()?;
-
-    Ok(m)
+#[pyo3::pymodule]
+pub(crate) mod dh {
+    #[pymodule_export]
+    use super::{
+        from_der_parameters, from_pem_parameters, generate_parameters, DHParameterNumbers,
+        DHParameters, DHPrivateKey, DHPrivateNumbers, DHPublicKey, DHPublicNumbers,
+    };
 }
diff --git a/src/rust/src/backend/dsa.rs b/src/rust/src/backend/dsa.rs
index d0218d45ba98..f46cb2860d33 100644
--- a/src/rust/src/backend/dsa.rs
+++ b/src/rust/src/backend/dsa.rs
@@ -7,7 +7,6 @@ use crate::buf::CffiBuf;
 use crate::error::{CryptographyError, CryptographyResult};
 use crate::exceptions;
 use pyo3::types::PyAnyMethods;
-use pyo3::types::PyModuleMethods;
 
 #[pyo3::pyclass(
     frozen,
@@ -499,18 +498,11 @@ impl DsaParameterNumbers {
     }
 }
 
-pub(crate) fn create_module(
-    py: pyo3::Python<'_>,
-) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::types::PyModule>> {
-    let m = pyo3::types::PyModule::new_bound(py, "dsa")?;
-    m.add_function(pyo3::wrap_pyfunction_bound!(generate_parameters, &m)?)?;
-
-    m.add_class::<DsaPrivateKey>()?;
-    m.add_class::<DsaPublicKey>()?;
-    m.add_class::<DsaParameters>()?;
-    m.add_class::<DsaPrivateNumbers>()?;
-    m.add_class::<DsaPublicNumbers>()?;
-    m.add_class::<DsaParameterNumbers>()?;
-
-    Ok(m)
+#[pyo3::pymodule]
+pub(crate) mod dsa {
+    #[pymodule_export]
+    use super::{
+        generate_parameters, DsaParameterNumbers, DsaParameters, DsaPrivateKey, DsaPrivateNumbers,
+        DsaPublicKey, DsaPublicNumbers,
+    };
 }
diff --git a/src/rust/src/backend/ec.rs b/src/rust/src/backend/ec.rs
index 6410add35cbe..15735458d3a1 100644
--- a/src/rust/src/backend/ec.rs
+++ b/src/rust/src/backend/ec.rs
@@ -5,7 +5,7 @@
 use std::collections::hash_map::DefaultHasher;
 use std::hash::{Hash, Hasher};
 
-use pyo3::types::{PyAnyMethods, PyDictMethods, PyModuleMethods};
+use pyo3::types::{PyAnyMethods, PyDictMethods};
 
 use crate::backend::utils;
 use crate::buf::CffiBuf;
@@ -670,19 +670,11 @@ impl EllipticCurvePublicNumbers {
     }
 }
 
-pub(crate) fn create_module(
-    py: pyo3::Python<'_>,
-) -> pyo3::PyResult<pyo3::Bound<'_, pyo3::types::PyModule>> {
-    let m = pyo3::types::PyModule::new_bound(py, "ec")?;
-    m.add_function(pyo3::wrap_pyfunction_bound!(curve_supported, &m)?)?;
-    m.add_function(pyo3::wrap_pyfunction_bound!(generate_private_key, &m)?)?;
-    m.add_function(pyo3::wrap_pyfunction_bound!(derive_private_key, &m)?)?;
-    m.add_function(pyo3::wrap_pyfunction_bound!(from_public_bytes, &m)?)?;
-
-    m.add_class::<ECPrivateKey>()?;
-    m.add_class::<ECPublicKey>()?;
-    m.add_class::<EllipticCurvePrivateNumbers>()?;
-    m.add_class::<EllipticCurvePublicNumbers>()?;
-
-    Ok(m)
+#[pyo3::pymodule]
+pub(crate) mod ec {
+    #[pymodule_export]
+    use super::{
+        curve_supported, derive_private_key, from_public_bytes, generate_private_key, ECPrivateKey,
+        ECPublicKey, EllipticCurvePrivateNumbers, EllipticCurvePublicNumbers,
+    };
 }
diff --git a/src/rust/src/backend/mod.rs b/src/rust/src/backend/mod.rs
index 4ee5f6bee124..0a36c4c70192 100644
--- a/src/rust/src/backend/mod.rs
+++ b/src/rust/src/backend/mod.rs
@@ -29,9 +29,6 @@ pub(crate) fn add_to_module(module: &pyo3::Bound<'_, pyo3::types::PyModule>) ->
     module.add_submodule(&aead::create_module(module.py())?)?;
     module.add_submodule(&ciphers::create_module(module.py())?)?;
     module.add_submodule(&cmac::create_module(module.py())?)?;
-    module.add_submodule(&dh::create_module(module.py())?)?;
-    module.add_submodule(&dsa::create_module(module.py())?)?;
-    module.add_submodule(&ec::create_module(module.py())?)?;
 
     module.add_submodule(&rsa::create_module(module.py())?)?;
 
diff --git a/src/rust/src/lib.rs b/src/rust/src/lib.rs
index 3f9568126f54..dbf497aa73a0 100644
--- a/src/rust/src/lib.rs
+++ b/src/rust/src/lib.rs
@@ -156,6 +156,12 @@ mod _rust {
         #[pymodule_export]
         use super::super::{is_fips_enabled, openssl_version, openssl_version_text};
         #[pymodule_export]
+        use crate::backend::dh::dh;
+        #[pymodule_export]
+        use crate::backend::dsa::dsa;
+        #[pymodule_export]
+        use crate::backend::ec::ec;
+        #[pymodule_export]
         use crate::backend::ed25519::ed25519;
         #[cfg(all(not(CRYPTOGRAPHY_IS_LIBRESSL), not(CRYPTOGRAPHY_IS_BORINGSSL)))]
         #[pymodule_export]