From 3f978c6723f5403702bb11b352f99f12ee856667 Mon Sep 17 00:00:00 2001
From: Bogdan Opanchuk <bogdan@opanchuk.net>
Date: Sun, 14 Mar 2021 21:11:36 -0700
Subject: [PATCH] Expose more types and serialization in Python bindings

---
 umbral-pre-python/src/lib.rs             | 82 +++++++++++++++++++++++-
 umbral-pre-python/umbral_pre/__init__.py |  3 +
 2 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/umbral-pre-python/src/lib.rs b/umbral-pre-python/src/lib.rs
index 670a15de..829a689c 100644
--- a/umbral-pre-python/src/lib.rs
+++ b/umbral-pre-python/src/lib.rs
@@ -20,9 +20,21 @@ impl SecretKey {
             backend: umbral_pre::SecretKey::random(),
         }
     }
+
+    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::SecretKey::from_encrypted_bytes(password, bytes)?;
+        Some(Self { backend: sk })
+    }
 }
 
 #[pyclass(module = "umbral")]
+#[derive(PartialEq)]
 pub struct PublicKey {
     backend: umbral_pre::PublicKey,
 }
@@ -35,14 +47,53 @@ impl PublicKey {
             backend: umbral_pre::PublicKey::from_secret_key(&sk.backend),
         }
     }
+
+    pub fn __bytes__(&self, py: Python) -> PyObject {
+        let serialized = self.backend.to_array();
+        PyBytes::new(py, serialized.as_slice()).into()
+    }
+
+    #[staticmethod]
+    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
+        let backend_pubkey = umbral_pre::PublicKey::from_bytes(bytes)?;
+        Some(Self {
+            backend: backend_pubkey,
+        })
+    }
+}
+
+#[pyproto]
+impl PyObjectProtocol for PublicKey {
+    fn __richcmp__(&self, other: PyRef<PublicKey>, op: CompareOp) -> PyResult<bool> {
+        match op {
+            CompareOp::Eq => Ok(self == &*other),
+            CompareOp::Ne => Ok(self != &*other),
+            _ => Err(PyTypeError::new_err("PublicKey objects are not ordered")),
+        }
+    }
 }
 
 #[pyclass(module = "umbral")]
-#[derive(Clone)]
 pub struct Capsule {
     backend: umbral_pre::Capsule,
 }
 
+#[pymethods]
+impl Capsule {
+    pub fn __bytes__(&self, py: Python) -> PyObject {
+        let serialized = self.backend.to_array();
+        PyBytes::new(py, serialized.as_slice()).into()
+    }
+
+    #[staticmethod]
+    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
+        let backend_capsule = umbral_pre::Capsule::from_bytes(bytes)?;
+        Some(Self {
+            backend: backend_capsule,
+        })
+    }
+}
+
 #[pyfunction]
 pub fn encrypt(py: Python, pk: &PublicKey, plaintext: &[u8]) -> (Capsule, PyObject) {
     let (capsule, ciphertext) = umbral_pre::encrypt(&pk.backend, plaintext).unwrap();
@@ -83,6 +134,19 @@ impl KeyFrag {
             receiving_pk.map(|pk| &pk.backend),
         )
     }
+
+    pub fn __bytes__(&self, py: Python) -> PyObject {
+        let serialized = self.backend.to_array();
+        PyBytes::new(py, serialized.as_slice()).into()
+    }
+
+    #[staticmethod]
+    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
+        let backend_kfrag = umbral_pre::KeyFrag::from_bytes(bytes)?;
+        Some(Self {
+            backend: backend_kfrag,
+        })
+    }
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -135,6 +199,19 @@ impl CapsuleFrag {
             &receiving_pk.backend,
         )
     }
+
+    pub fn __bytes__(&self, py: Python) -> PyObject {
+        let serialized = self.backend.to_array();
+        PyBytes::new(py, serialized.as_slice()).into()
+    }
+
+    #[staticmethod]
+    pub fn from_bytes(bytes: &[u8]) -> Option<Self> {
+        let backend_cfrag = umbral_pre::CapsuleFrag::from_bytes(bytes)?;
+        Some(Self {
+            backend: backend_cfrag,
+        })
+    }
 }
 
 #[pyfunction]
@@ -174,6 +251,9 @@ pub fn decrypt_reencrypted(
 fn _umbral(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<SecretKey>()?;
     m.add_class::<PublicKey>()?;
+    m.add_class::<Capsule>()?;
+    m.add_class::<KeyFrag>()?;
+    m.add_class::<CapsuleFrag>()?;
     m.add_function(wrap_pyfunction!(encrypt, m)?).unwrap();
     m.add_function(wrap_pyfunction!(decrypt_original, m)?)
         .unwrap();
diff --git a/umbral-pre-python/umbral_pre/__init__.py b/umbral-pre-python/umbral_pre/__init__.py
index b9b20d0c..a2ed40a8 100644
--- a/umbral-pre-python/umbral_pre/__init__.py
+++ b/umbral-pre-python/umbral_pre/__init__.py
@@ -1,6 +1,9 @@
 from ._umbral import (
     SecretKey,
     PublicKey,
+    Capsule,
+    KeyFrag,
+    CapsuleFrag,
     encrypt,
     decrypt_original,
     decrypt_reencrypted,