Add PkEncryption support

Author: MatMaul

Cargo.toml

@@ -7,10 +7,14 @@ edition = "2021"
 name = "vodozemac"
 crate-type = ["cdylib"]
 
+[features]
+default = ["insecure-pk-encryption"]
+insecure-pk-encryption = ["vodozemac/insecure-pk-encryption"]
+
 [dependencies]
 paste = "1.0.15"
 thiserror = "1.0.63"
-vodozemac = "0.7.0"
+vodozemac = { git = "https://github.com/matrix-org/vodozemac.git", branch = "poljar/pk-dekurcina" }
 
 [package.metadata.maturin]
 name = "vodozemac"

src/error.rs

@@ -145,3 +145,26 @@ impl From<PickleError> for PyErr {
         PickleException::new_err(e.to_string())
     }
 }
+
+#[cfg(feature = "insecure-pk-encryption")]
+#[derive(Debug, Error)]
+pub enum PkEncryptionError {
+    #[error("The key doesn't have the correct size, got {0}, expected 32 bytes")]
+    InvalidKeySize(usize),
+    #[error(transparent)]
+    Decode(#[from] vodozemac::pk_encryption::Error),
+}
+
+#[cfg(feature = "insecure-pk-encryption")]
+pyo3::create_exception!(module, PkInvalidKeySizeException, pyo3::exceptions::PyValueError);
+pyo3::create_exception!(module, PkDecodeException, pyo3::exceptions::PyValueError);
+
+#[cfg(feature = "insecure-pk-encryption")]
+impl From<PkEncryptionError> for PyErr {
+    fn from(e: PkEncryptionError) -> Self {
+        match e {
+            PkEncryptionError::InvalidKeySize(_) => PkInvalidKeySizeException::new_err(e.to_string()),
+            PkEncryptionError::Decode(_) => PkDecodeException::new_err(e.to_string()),
+        }
+    }
+}

src/lib.rs

@@ -1,6 +1,8 @@
 mod account;
 mod error;
 mod group_sessions;
+#[cfg(feature = "insecure-pk-encryption")]
+mod pk_encryption;
 mod sas;
 mod session;
 mod types;
@@ -35,6 +37,12 @@ fn my_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<sas::Sas>()?;
     m.add_class::<group_sessions::GroupSession>()?;
     m.add_class::<group_sessions::InboundGroupSession>()?;
+    #[cfg(feature = "insecure-pk-encryption")]
+    m.add_class::<pk_encryption::PkDecryption>()?;
+    #[cfg(feature = "insecure-pk-encryption")]
+    m.add_class::<pk_encryption::PkEncryption>()?;
+    #[cfg(feature = "insecure-pk-encryption")]
+    m.add_class::<pk_encryption::Message>()?;
 
     m.add_class::<types::Ed25519PublicKey>()?;
     m.add_class::<types::Curve25519PublicKey>()?;
@@ -63,6 +71,16 @@ fn my_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
         "MegolmDecryptionException",
         py.get_type_bound::<MegolmDecryptionException>(),
     )?;
+    #[cfg(feature = "insecure-pk-encryption")]
+    m.add(
+        "PkInvalidKeySizeException",
+        py.get_type_bound::<PkInvalidKeySizeException>(),
+    )?;
+    #[cfg(feature = "insecure-pk-encryption")]
+    m.add(
+        "PkDecodeException",
+        py.get_type_bound::<PkDecodeException>(),
+    )?;
 
     Ok(())
 }

src/pk_encryption.rs

@@ -0,0 +1,115 @@
+use pyo3::{
+    pyclass, pymethods,
+    types::{PyBytes, PyType},
+    Bound, Py, Python,
+};
+use vodozemac::Curve25519PublicKey;
+
+use crate::PkEncryptionError;
+
+#[pyclass]
+pub struct Message {
+    ciphertext: Vec<u8>,
+    mac: Vec<u8>,
+    ephemeral_key: Vec<u8>,
+}
+
+/// The decryption component of the PkEncryption support.
+#[pyclass]
+pub struct PkDecryption {
+    inner: vodozemac::pk_encryption::PkDecryption,
+}
+
+#[pymethods]
+impl PkDecryption {
+    /// Create a new random PkDecryption object.
+    #[new]
+    fn new() -> Self {
+        Self {
+            inner: vodozemac::pk_encryption::PkDecryption::new(),
+        }
+    }
+
+    /// Create a PkDecryption object from the secret key bytes.
+    #[classmethod]
+    fn from_bytes(_cls: &Bound<'_, PyType>, bytes: &[u8]) -> Result<Self, PkEncryptionError> {
+        let key: &[u8; 32] = bytes
+            .try_into()
+            .map_err(|_| PkEncryptionError::InvalidKeySize(bytes.len()))?;
+
+        Ok(Self {
+            inner: vodozemac::pk_encryption::PkDecryption::from_bytes(&key),
+        })
+    }
+
+    // The secret key used to decrypt messages.
+    #[getter]
+    pub fn key(&self) -> Py<PyBytes> {
+        Python::with_gil(|py| PyBytes::new_bound(py, self.inner.to_bytes().as_slice()).into())
+    }
+
+    // The public key used to encrypt messages for this decryption object.
+    #[getter]
+    pub fn public_key(&self) -> Py<PyBytes> {
+        Python::with_gil(|py| PyBytes::new_bound(py, self.inner.public_key().as_bytes()).into())
+    }
+
+    /// Decrypt a ciphertext. See the PkEncryption::encrypt function
+    /// for descriptions of the ephemeral_key and mac arguments.
+    pub fn decrypt(&self, message: &Message) -> Result<Py<PyBytes>, PkEncryptionError> {
+        let ephemeral_key_bytes: [u8; 32] = message
+            .ephemeral_key
+            .as_slice()
+            .try_into()
+            .map_err(|_| PkEncryptionError::InvalidKeySize(message.ephemeral_key.len()))?;
+
+        let message = vodozemac::pk_encryption::Message {
+            ciphertext: message.ciphertext.clone(),
+            mac: message.mac.clone(),
+            ephemeral_key: Curve25519PublicKey::from_bytes(ephemeral_key_bytes),
+        };
+
+        self.inner
+            .decrypt(&message)
+            .map(|vec| Python::with_gil(|py| PyBytes::new_bound(py, vec.as_slice()).into()))
+            .map_err(|e| PkEncryptionError::Decode(e))
+    }
+}
+
+/// The encryption component of PkEncryption support.
+#[pyclass]
+pub struct PkEncryption {
+    inner: vodozemac::pk_encryption::PkEncryption,
+}
+
+#[pymethods]
+impl PkEncryption {
+    /// Create a new PkEncryption object from public key bytes.
+    #[classmethod]
+    fn from_public_key(_cls: &Bound<'_, PyType>, bytes: &[u8]) -> Result<Self, PkEncryptionError> {
+        let key: &[u8; 32] = bytes
+            .try_into()
+            .map_err(|_| PkEncryptionError::InvalidKeySize(bytes.len()))?;
+
+        Ok(Self {
+            inner: vodozemac::pk_encryption::PkEncryption::from_key(
+                Curve25519PublicKey::from_bytes(*key),
+            ),
+        })
+    }
+
+    /// Encrypt a plaintext for the recipient. Writes to the ciphertext, mac, and
+    /// ephemeral_key buffers, whose values should be sent to the recipient. mac is
+    /// a Message Authentication Code to ensure that the data is received and
+    /// decrypted properly. ephemeral_key is the public part of the ephemeral key
+    /// used (together with the recipient's key) to generate a symmetric encryption
+    /// key.
+    pub fn encrypt(&self, message: &[u8]) -> Message {
+        let msg = self.inner.encrypt(message);
+        Message {
+            ciphertext: msg.ciphertext.to_vec(),
+            mac: msg.mac.to_vec(),
+            ephemeral_key: msg.ephemeral_key.to_vec(),
+        }
+    }
+}

tests/pk_encryption_test.py

@@ -0,0 +1,32 @@
+import importlib
+import pytest
+
+SKIP = False
+try:
+    from vodozemac import PkEncryption, PkDecryption, PkDecodeException, PkInvalidKeySizeException
+except ImportError:
+    SKIP = True
+
+CLEARTEXT = b"test"
+
+@pytest.mark.skipif(SKIP, reason="Can't import, feature probably disabled")
+class TestClass(object):
+    def test_encrypt_decrypt(self):
+        d = PkDecryption()
+        e = PkEncryption.from_public_key(d.public_key)
+
+        decoded = d.decrypt(e.encrypt(CLEARTEXT))
+        assert decoded == CLEARTEXT
+
+    def test_encrypt_decrypt_with_wrong_key(self):
+        wrong_e = PkEncryption.from_public_key(PkDecryption().public_key)
+        with pytest.raises(PkDecodeException, match="MAC tag mismatch"):
+            PkDecryption().decrypt(wrong_e.encrypt(CLEARTEXT))
+
+    def test_encrypt_key_wrong_size(self):
+        with pytest.raises(PkInvalidKeySizeException, match="expected 32 bytes"):
+            PkEncryption.from_public_key(bytes())
+
+    def test_keys_wrong_size(self):
+        with pytest.raises(PkInvalidKeySizeException, match="expected 32 bytes"):
+            PkDecryption.from_bytes(bytes())