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 { + 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 { + 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 { + let backend_pubkey = umbral_pre::PublicKey::from_bytes(bytes)?; + Some(Self { + backend: backend_pubkey, + }) + } +} + +#[pyproto] +impl PyObjectProtocol for PublicKey { + fn __richcmp__(&self, other: PyRef, op: CompareOp) -> PyResult { + 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 { + 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 { + 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 { + 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::()?; m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; 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,