From ec3265557095638a3f5d662a1fa73070ada38ce9 Mon Sep 17 00:00:00 2001
From: Bogdan Opanchuk <bogdan@opanchuk.net>
Date: Sun, 21 Mar 2021 15:05:31 -0700
Subject: [PATCH] py-bindings: expose SecretKeyFactory

---
 umbral-pre-python/src/lib.rs             | 46 ++++++++++++++++++++++++
 umbral-pre-python/umbral_pre/__init__.py |  1 +
 2 files changed, 47 insertions(+)

diff --git a/umbral-pre-python/src/lib.rs b/umbral-pre-python/src/lib.rs
index af83edf3..2e646b12 100644
--- a/umbral-pre-python/src/lib.rs
+++ b/umbral-pre-python/src/lib.rs
@@ -8,6 +8,7 @@ use pyo3::PyObjectProtocol;
 use umbral_pre::SerializableToArray;
 
 #[pyclass(module = "umbral")]
+#[derive(PartialEq)]
 pub struct SecretKey {
     backend: umbral_pre::SecretKey,
 }
@@ -33,6 +34,50 @@ impl SecretKey {
     }
 }
 
+#[pyproto]
+impl PyObjectProtocol for SecretKey {
+    fn __richcmp__(&self, other: PyRef<SecretKey>, op: CompareOp) -> PyResult<bool> {
+        match op {
+            CompareOp::Eq => Ok(self == &*other),
+            CompareOp::Ne => Ok(self != &*other),
+            _ => Err(PyTypeError::new_err("SecretKey objects are not ordered")),
+        }
+    }
+}
+
+#[pyclass(module = "umbral")]
+pub struct SecretKeyFactory {
+    backend: umbral_pre::SecretKeyFactory,
+}
+
+#[pymethods]
+impl SecretKeyFactory {
+    #[staticmethod]
+    pub fn random() -> Self {
+        Self {
+            backend: umbral_pre::SecretKeyFactory::random(),
+        }
+    }
+
+    pub fn secret_key_by_label(&self, label: &[u8]) -> Option<SecretKey> {
+        let backend_sk = self.backend.secret_key_by_label(label)?;
+        Some(SecretKey {
+            backend: backend_sk,
+        })
+    }
+
+    pub fn to_encrypted_bytes(&self, py: Python, password: &[u8]) -> Option<PyObject> {
+        let ciphertext = self.backend.to_encrypted_bytes(password)?;
+        Some(PyBytes::new(py, ciphertext.as_slice()).into())
+    }
+
+    #[staticmethod]
+    pub fn from_encrypted_bytes(password: &[u8], bytes: &[u8]) -> Option<Self> {
+        let sk = umbral_pre::SecretKeyFactory::from_encrypted_bytes(password, bytes)?;
+        Some(Self { backend: sk })
+    }
+}
+
 #[pyclass(module = "umbral")]
 #[derive(PartialEq)]
 pub struct PublicKey {
@@ -252,6 +297,7 @@ pub fn decrypt_reencrypted(
 #[pymodule]
 fn _umbral(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<SecretKey>()?;
+    m.add_class::<SecretKeyFactory>()?;
     m.add_class::<PublicKey>()?;
     m.add_class::<Capsule>()?;
     m.add_class::<KeyFrag>()?;
diff --git a/umbral-pre-python/umbral_pre/__init__.py b/umbral-pre-python/umbral_pre/__init__.py
index a2ed40a8..4b2a39c3 100644
--- a/umbral-pre-python/umbral_pre/__init__.py
+++ b/umbral-pre-python/umbral_pre/__init__.py
@@ -1,5 +1,6 @@
 from ._umbral import (
     SecretKey,
+    SecretKeyFactory,
     PublicKey,
     Capsule,
     KeyFrag,