Add Python Attribute Accessors and Cross-Library Compatibility for PkEncryption

requirements-dev.txt

@@ -1,2 +1,3 @@
 maturin
 pytest>=4.0
+python-olm==3.2.16

src/pk_encryption.rs

@@ -1,7 +1,7 @@
 use pyo3::{
     pyclass, pymethods,
-    types::{PyBytes, PyType},
-    Bound, Py, Python,
+    types::{PyBytes, PyString, PyType},
+    Bound, IntoPyObject, Py, PyResult, Python,
 };
 
 use crate::{
@@ -13,17 +13,84 @@ use crate::{
 #[pyclass]
 pub struct Message {
     /// The ciphertext of the message.
+    #[pyo3(get)]
     ciphertext: Vec<u8>,
     /// The message authentication code of the message.
     ///
     /// *Warning*: As stated in the module description, this does not
     /// authenticate the message.
+    #[pyo3(get)]
     mac: Vec<u8>,
     /// The ephemeral Curve25519PublicKey of the message which was used to
     /// derive the individual message key.
+    #[pyo3(get)]
     ephemeral_key: Vec<u8>,
 }
 
+#[pymethods]
+impl Message {
+    /// Create a new Message object from its components.
+    ///
+    /// This constructor creates a Message object that represents an encrypted message
+    /// using the `m.megolm_backup.v1.curve25519-aes-sha2` algorithm.
+    ///
+    /// # Arguments
+    /// * `ciphertext` - The encrypted content of the message
+    /// * `mac` - The message authentication code
+    /// * `ephemeral_key` - The ephemeral public key used during encryption
+    #[new]
+    fn new(ciphertext: Vec<u8>, mac: Vec<u8>, ephemeral_key: Vec<u8>) -> Self {
+        Message { ciphertext, mac, ephemeral_key }
+    }
+
+    /// Create a new Message object from unpadded Base64-encoded components.
+    ///
+    /// This function decodes the given Base64 strings and returns a `Message`
+    /// with the resulting byte vectors.
+    ///
+    /// # Arguments
+    /// * `ciphertext` - Unpadded Base64-encoded ciphertext
+    /// * `mac` - Unpadded Base64-encoded message authentication code
+    /// * `ephemeral_key` - Unpadded Base64-encoded ephemeral key
+    #[classmethod]
+    fn from_base64(
+        _cls: &Bound<'_, PyType>,
+        ciphertext: &str,
+        mac: &str,
+        ephemeral_key: &str,
+    ) -> Self {
+        let decoded_ciphertext =
+            vodozemac::base64_decode(ciphertext).expect("Failed to decode ciphertext");
+        let decoded_mac = vodozemac::base64_decode(mac).expect("Failed to decode mac");
+        let decoded_ephemeral_key =
+            vodozemac::base64_decode(ephemeral_key).expect("Failed to decode ephemeral_key");
+
+        Self {
+            ciphertext: decoded_ciphertext,
+            mac: decoded_mac,
+            ephemeral_key: decoded_ephemeral_key,
+        }
+    }
+
+    /// Convert the message components to unpadded Base64-encoded strings.
+    ///
+    /// Returns a tuple of (ciphertext, mac, ephemeral_key) as unpadded Base64 strings.
+    fn to_base64<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<(Bound<'py, PyString>, Bound<'py, PyString>, Bound<'py, PyString>)> {
+        let ciphertext_b64 = vodozemac::base64_encode(&self.ciphertext);
+        let mac_b64 = vodozemac::base64_encode(&self.mac);
+        let ephemeral_key_b64 = vodozemac::base64_encode(&self.ephemeral_key);
+
+        Ok((
+            ephemeral_key_b64.into_pyobject(py)?,
+            mac_b64.into_pyobject(py)?,
+            ciphertext_b64.into_pyobject(py)?,
+        ))
+    }
+}
+
 /// ☣️  Compat support for libolm's PkDecryption.
 ///
 /// This implements the `m.megolm_backup.v1.curve25519-aes-sha2` described in

tests/pk_encryption_test.py

@@ -1,10 +1,20 @@
 import importlib
 import pytest
+import olm
+import base64
 
-from vodozemac import Curve25519SecretKey, Curve25519PublicKey, PkEncryption, PkDecryption, PkDecodeException
+from vodozemac import (
+    Curve25519SecretKey,
+    Curve25519PublicKey,
+    PkEncryption,
+    PkDecryption,
+    PkDecodeException,
+    Message,
+)
 
 CLEARTEXT = b"test"
 
+
 class TestClass(object):
     def test_encrypt_decrypt(self):
         d = PkDecryption()
@@ -28,3 +38,50 @@ def test_encrypt_decrypt_with_serialized_keys(self):
 
         decoded = d.decrypt(e.encrypt(CLEARTEXT))
         assert decoded == CLEARTEXT
+
+    def test_encrypt_message_attr(self):
+        """Test that the Message object has accessible Python attributes (mac, ciphertext, ephemeral_key)."""
+        decryption = PkDecryption()
+        encryption = PkEncryption.from_key(decryption.public_key)
+
+        message = encryption.encrypt(CLEARTEXT)
+
+        assert message.mac is not None
+        assert message.ciphertext is not None
+        assert message.ephemeral_key is not None
+
+    def test_olm_encrypt_vodo_decrypt(self):
+        """Test encrypting with Olm and decrypting with Vodo."""
+        vodo_decryption = PkDecryption()
+        olm_encrypts = olm.pk.PkEncryption(vodo_decryption.public_key.to_base64())
+        olm_msg = olm_encrypts.encrypt(CLEARTEXT)
+
+        vodo_msg = Message.from_base64(
+            olm_msg.ciphertext,
+            olm_msg.mac,
+            olm_msg.ephemeral_key,
+        )
+
+        # Decrypt the message with Vodo
+        decrypted_plaintext = vodo_decryption.decrypt(vodo_msg)
+        assert decrypted_plaintext == CLEARTEXT
+
+    def test_vodo_encrypt_olm_decrypt(self):
+        """Test encrypting with Vodo and decrypting with Olm."""
+        olm_decryption = olm.pk.PkDecryption()
+
+        public_key = Curve25519PublicKey.from_base64(olm_decryption.public_key)
+        vodo_encryption = PkEncryption.from_key(public_key)
+        vodo_msg = vodo_encryption.encrypt(CLEARTEXT)
+
+        ephemeral_key_b64, mac_b64, ciphertext_b64 = vodo_msg.to_base64()
+
+        olm_msg = olm.pk.PkMessage(
+            ephemeral_key_b64,
+            mac_b64,
+            ciphertext_b64
+        )
+
+        # Decrypt the message with Olm
+        decrypted_plaintext = olm_decryption.decrypt(olm_msg)
+        assert decrypted_plaintext.encode("utf-8") == CLEARTEXT