diff --git a/requirements-dev.txt b/requirements-dev.txt index 20faf86..0846534 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,2 +1,3 @@ maturin pytest>=4.0 +python-olm==3.2.16 diff --git a/src/pk_encryption.rs b/src/pk_encryption.rs index f93f862..3ffa37f 100644 --- a/src/pk_encryption.rs +++ b/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, /// The message authentication code of the message. /// /// *Warning*: As stated in the module description, this does not /// authenticate the message. + #[pyo3(get)] mac: Vec, /// The ephemeral Curve25519PublicKey of the message which was used to /// derive the individual message key. + #[pyo3(get)] ephemeral_key: Vec, } +#[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, mac: Vec, ephemeral_key: Vec) -> 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 diff --git a/tests/pk_encryption_test.py b/tests/pk_encryption_test.py index 39a093d..974e922 100644 --- a/tests/pk_encryption_test.py +++ b/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