Add PkEncryption support

Cargo.toml

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

src/error.rs

@@ -139,3 +139,29 @@ impl From<PickleError> for PyErr {
         PickleException::new_err(e.to_string())
     }
 }
+
+#[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),
+}
+
+pyo3::create_exception!(
+    module,
+    PkInvalidKeySizeException,
+    pyo3::exceptions::PyValueError
+);
+pyo3::create_exception!(module, PkDecodeException, pyo3::exceptions::PyValueError);
+
+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,7 @@
 mod account;
 mod error;
 mod group_sessions;
+mod pk_encryption;
 mod sas;
 mod session;
 mod types;
@@ -26,6 +27,11 @@ fn my_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<types::Ed25519PublicKey>()?;
     m.add_class::<types::Ed25519Signature>()?;
     m.add_class::<types::Curve25519PublicKey>()?;
+    m.add_class::<types::Curve25519SecretKey>()?;
+
+    m.add_class::<pk_encryption::PkDecryption>()?;
+    m.add_class::<pk_encryption::PkEncryption>()?;
+    m.add_class::<pk_encryption::Message>()?;
 
     m.add("KeyException", py.get_type_bound::<KeyException>())?;
     m.add(
@@ -55,6 +61,14 @@ fn my_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
         "MegolmDecryptionException",
         py.get_type_bound::<MegolmDecryptionException>(),
     )?;
+    m.add(
+        "PkInvalidKeySizeException",
+        py.get_type_bound::<PkInvalidKeySizeException>(),
+    )?;
+    m.add(
+        "PkDecodeException",
+        py.get_type_bound::<PkDecodeException>(),
+    )?;
 
     Ok(())
 }

src/pk_encryption.rs

@@ -0,0 +1,139 @@
+use pyo3::{
+    pyclass, pymethods,
+    types::{PyBytes, PyType},
+    Bound, Py, Python,
+};
+
+use crate::{
+    types::{Curve25519PublicKey, Curve25519SecretKey},
+    PkEncryptionError,
+};
+
+/// A message that was encrypted using a PkEncryption object.
+#[pyclass]
+pub struct Message {
+    /// The ciphertext of the message.
+    ciphertext: Vec<u8>,
+    /// The message authentication code of the message.
+    ///
+    /// *Warning*: As stated in the module description, this does not
+    /// authenticate the message.
+    mac: Vec<u8>,
+    /// The ephemeral Curve25519PublicKey of the message which was used to
+    /// derive the individual message key.
+    ephemeral_key: Vec<u8>,
+}
+
+/// ☣️  Compat support for libolm's PkDecryption.
+///
+/// This implements the `m.megolm_backup.v1.curve25519-aes-sha2` described in
+/// the Matrix [spec]. This is a asymmetric encryption scheme based on
+/// Curve25519.
+///
+/// **Warning**: Please note the algorithm contains a critical flaw and does not
+/// provide authentication of the ciphertext.
+///
+/// [spec]: https://spec.matrix.org/v1.11/client-server-api/#backup-algorithm-mmegolm_backupv1curve25519-aes-sha2
+#[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_key(
+        _cls: &Bound<'_, PyType>,
+        key: Curve25519SecretKey,
+    ) -> Result<Self, PkEncryptionError> {
+        Ok(Self {
+            inner: vodozemac::pk_encryption::PkDecryption::from_key(key.inner),
+        })
+    }
+
+    /// The secret key used to decrypt messages.
+    #[getter]
+    pub fn key(&self) -> Curve25519SecretKey {
+        Curve25519SecretKey::from(self.inner.secret_key().clone())
+    }
+
+    /// The public key used to encrypt messages for this decryption object.
+    #[getter]
+    pub fn public_key(&self) -> Curve25519PublicKey {
+        Curve25519PublicKey::from(self.inner.public_key())
+    }
+
+    /// 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: vodozemac::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))
+    }
+}
+
+/// ☣️  Compat support for libolm's PkEncryption.
+///
+/// This implements the `m.megolm_backup.v1.curve25519-aes-sha2` described in
+/// the Matrix [spec]. This is a asymmetric encryption scheme based on
+/// Curve25519.
+///
+/// **Warning**: Please note the algorithm contains a critical flaw and does not
+/// provide authentication of the ciphertext.
+///
+/// [spec]: https://spec.matrix.org/v1.11/client-server-api/#backup-algorithm-mmegolm_backupv1curve25519-aes-sha2
+#[pyclass]
+pub struct PkEncryption {
+    inner: vodozemac::pk_encryption::PkEncryption,
+}
+
+#[pymethods]
+impl PkEncryption {
+    /// Create a new PkEncryption object from public key.
+    #[classmethod]
+    fn from_key(
+        _cls: &Bound<'_, PyType>,
+        key: Curve25519PublicKey,
+    ) -> Result<Self, PkEncryptionError> {
+        Ok(Self {
+            inner: vodozemac::pk_encryption::PkEncryption::from_key(key.inner),
+        })
+    }
+
+    /// 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(),
+        }
+    }
+}

src/types/curve25519.rs

@@ -1,5 +1,9 @@
-use crate::error::*;
-use pyo3::{prelude::*, types::PyType};
+use crate::{convert_to_pybytes, error::*};
+use pyo3::{
+    prelude::*,
+    types::{PyBytes, PyType},
+};
+use vodozemac::{base64_decode, base64_encode};
 
 #[pyclass]
 #[derive(Clone)]
@@ -22,14 +26,101 @@ impl Curve25519PublicKey {
         })
     }
 
+    #[classmethod]
+    pub fn from_bytes(_cls: &Bound<'_, PyType>, bytes: &[u8]) -> Result<Self, KeyError> {
+        let key: &[u8; 32] = bytes.try_into().map_err(|_| {
+            KeyError::from(vodozemac::KeyError::InvalidKeyLength {
+                key_type: "Curve25519PublicKey",
+                expected_length: 32,
+                length: bytes.len(),
+            })
+        })?;
+
+        Ok(Self {
+            inner: vodozemac::Curve25519PublicKey::from_slice(key)?,
+        })
+    }
+
     pub fn to_base64(&self) -> String {
         self.inner.to_base64()
     }
 
+    pub fn to_bytes(&self) -> Py<PyBytes> {
+        convert_to_pybytes(self.inner.to_bytes())
+    }
+
     #[classattr]
     const __hash__: Option<PyObject> = None;
 
     fn __eq__(&self, other: &Self) -> bool {
         self.inner == other.inner
     }
 }
+
+#[pyclass]
+#[derive(Clone)]
+pub struct Curve25519SecretKey {
+    pub(crate) inner: vodozemac::Curve25519SecretKey,
+}
+
+impl From<vodozemac::Curve25519SecretKey> for Curve25519SecretKey {
+    fn from(value: vodozemac::Curve25519SecretKey) -> Self {
+        Self { inner: value }
+    }
+}
+
+#[pymethods]
+impl Curve25519SecretKey {
+    #[new]
+    fn new() -> Self {
+        Self {
+            inner: vodozemac::Curve25519SecretKey::new(),
+        }
+    }
+
+    #[classmethod]
+    pub fn from_base64(_cls: &Bound<'_, PyType>, key: &str) -> Result<Self, KeyError> {
+        Self::from_bytes(
+            _cls,
+            base64_decode(key)
+                .map_err(|e| KeyError::from(vodozemac::KeyError::Base64Error(e)))?
+                .as_slice(),
+        )
+    }
+
+    #[classmethod]
+    pub fn from_bytes(_cls: &Bound<'_, PyType>, bytes: &[u8]) -> Result<Self, KeyError> {
+        let key: &[u8; 32] = bytes.try_into().map_err(|_| {
+            KeyError::from(vodozemac::KeyError::InvalidKeyLength {
+                key_type: "Curve25519SecretKey",
+                expected_length: 32,
+                length: bytes.len(),
+            })
+        })?;
+
+        Ok(Self {
+            inner: vodozemac::Curve25519SecretKey::from_slice(key),
+        })
+    }
+
+    pub fn to_base64(&self) -> String {
+        base64_encode(self.inner.to_bytes().as_slice())
+    }
+
+    pub fn to_bytes(&self) -> Py<PyBytes> {
+        convert_to_pybytes(self.inner.to_bytes().as_slice())
+    }
+
+    pub fn public_key(&self) -> Curve25519PublicKey {
+        Curve25519PublicKey {
+            inner: vodozemac::Curve25519PublicKey::from(&self.inner),
+        }
+    }
+
+    #[classattr]
+    const __hash__: Option<PyObject> = None;
+
+    fn __eq__(&self, other: &Self) -> bool {
+        self.inner.to_bytes() == other.inner.to_bytes()
+    }
+}

tests/pk_encryption_test.py

@@ -0,0 +1,30 @@
+import importlib
+import pytest
+
+from vodozemac import Curve25519SecretKey, Curve25519PublicKey, PkEncryption, PkDecryption, PkDecodeException
+
+CLEARTEXT = b"test"
+
+class TestClass(object):
+    def test_encrypt_decrypt(self):
+        d = PkDecryption()
+        e = PkEncryption.from_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_key(PkDecryption().public_key)
+        with pytest.raises(PkDecodeException, match="MAC tag mismatch"):
+            PkDecryption().decrypt(wrong_e.encrypt(CLEARTEXT))
+
+    def test_encrypt_decrypt_with_serialized_keys(self):
+        secret_key = Curve25519SecretKey()
+        secret_key_bytes = secret_key.to_bytes()
+        public_key_bytes = secret_key.public_key().to_bytes()
+
+        d = PkDecryption.from_key(Curve25519SecretKey.from_bytes(secret_key_bytes))
+        e = PkEncryption.from_key(Curve25519PublicKey.from_bytes(public_key_bytes))
+
+        decoded = d.decrypt(e.encrypt(CLEARTEXT))
+        assert decoded == CLEARTEXT