Skip to content

Commit f9ebdfe

Browse files
TrevisGordanpoljarPaarthShah
authored
Add Python Attribute Accessors for PkEncryption (#15)
* Adds missing rust message attribute accessibility and test * Adds Message impl (pymethods) to enable object creation from python. Adds python-olm Dependencie (at 3.2.16). Adds PkEncryption (python-olm) compatibility tests. --------- Co-authored-by: Damir Jelić <poljar@termina.org.uk> Co-authored-by: Paarth Shah <mail@shahpaarth.com>
1 parent 288157e commit f9ebdfe

File tree

3 files changed

+112
-3
lines changed

3 files changed

+112
-3
lines changed

src/error.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,8 @@ pub enum PkEncryptionError {
136136
InvalidKeySize(usize),
137137
#[error(transparent)]
138138
Decode(#[from] vodozemac::pk_encryption::Error),
139+
#[error(transparent)]
140+
Mac(#[from] vodozemac::Base64DecodeError),
139141
}
140142

141143
pyo3::create_exception!(module, PkInvalidKeySizeException, pyo3::exceptions::PyValueError);
@@ -148,6 +150,7 @@ impl From<PkEncryptionError> for PyErr {
148150
PkInvalidKeySizeException::new_err(e.to_string())
149151
}
150152
PkEncryptionError::Decode(_) => PkDecodeException::new_err(e.to_string()),
153+
PkEncryptionError::Mac(_) => PkDecodeException::new_err(e.to_string()),
151154
}
152155
}
153156
}

src/pk_encryption.rs

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,17 +13,77 @@ use crate::{
1313
#[pyclass]
1414
pub struct Message {
1515
/// The ciphertext of the message.
16+
#[pyo3(get)]
1617
ciphertext: Vec<u8>,
1718
/// The message authentication code of the message.
1819
///
1920
/// *Warning*: As stated in the module description, this does not
2021
/// authenticate the message.
22+
#[pyo3(get)]
2123
mac: Vec<u8>,
2224
/// The ephemeral Curve25519PublicKey of the message which was used to
2325
/// derive the individual message key.
26+
#[pyo3(get)]
2427
ephemeral_key: Vec<u8>,
2528
}
2629

30+
#[pymethods]
31+
impl Message {
32+
/// Create a new Message object from its components.
33+
///
34+
/// This constructor creates a Message object that represents an encrypted
35+
/// message using the `m.megolm_backup.v1.curve25519-aes-sha2`
36+
/// algorithm.
37+
///
38+
/// # Arguments
39+
/// * `ciphertext` - The encrypted content of the message
40+
/// * `mac` - The message authentication code
41+
/// * `ephemeral_key` - The ephemeral public key used during encryption
42+
#[new]
43+
fn new(ciphertext: Vec<u8>, mac: Vec<u8>, ephemeral_key: Vec<u8>) -> Self {
44+
Message { ciphertext, mac, ephemeral_key }
45+
}
46+
47+
/// Create a new Message object from unpadded Base64-encoded components.
48+
///
49+
/// This function decodes the given Base64 strings and returns a `Message`
50+
/// with the resulting byte vectors.
51+
///
52+
/// # Arguments
53+
/// * `ciphertext` - Unpadded Base64-encoded ciphertext
54+
/// * `mac` - Unpadded Base64-encoded message authentication code
55+
/// * `ephemeral_key` - Unpadded Base64-encoded ephemeral key
56+
#[classmethod]
57+
fn from_base64(
58+
_cls: &Bound<'_, PyType>,
59+
ciphertext: &str,
60+
mac: &str,
61+
ephemeral_key: &str,
62+
) -> Result<Self, PkEncryptionError> {
63+
let decoded_ciphertext = vodozemac::base64_decode(ciphertext)?;
64+
let decoded_mac = vodozemac::base64_decode(mac)?;
65+
let decoded_ephemeral_key = vodozemac::base64_decode(ephemeral_key)?;
66+
67+
Ok(Self {
68+
ciphertext: decoded_ciphertext,
69+
mac: decoded_mac,
70+
ephemeral_key: decoded_ephemeral_key,
71+
})
72+
}
73+
74+
/// Convert the message components to unpadded Base64-encoded strings.
75+
///
76+
/// Returns a tuple of (ciphertext, mac, ephemeral_key) as unpadded Base64
77+
/// strings.
78+
fn to_base64(&self) -> Result<(String, String, String), PkEncryptionError> {
79+
let ciphertext_b64 = vodozemac::base64_encode(&self.ciphertext);
80+
let mac_b64 = vodozemac::base64_encode(&self.mac);
81+
let ephemeral_key_b64 = vodozemac::base64_encode(&self.ephemeral_key);
82+
83+
Ok((ephemeral_key_b64, mac_b64, ciphertext_b64))
84+
}
85+
}
86+
2787
/// ☣️ Compat support for libolm's PkDecryption.
2888
///
2989
/// This implements the `m.megolm_backup.v1.curve25519-aes-sha2` described in

tests/pk_encryption_test.py

Lines changed: 49 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,18 @@
1-
import importlib
2-
import pytest
1+
import base64
32

4-
from vodozemac import Curve25519SecretKey, Curve25519PublicKey, PkEncryption, PkDecryption, PkDecodeException
3+
import pytest
4+
from vodozemac import (
5+
Curve25519PublicKey,
6+
Curve25519SecretKey,
7+
Message,
8+
PkDecodeException,
9+
PkDecryption,
10+
PkEncryption,
11+
)
512

613
CLEARTEXT = b"test"
714

15+
816
class TestClass(object):
917
def test_encrypt_decrypt(self):
1018
d = PkDecryption()
@@ -28,3 +36,41 @@ def test_encrypt_decrypt_with_serialized_keys(self):
2836

2937
decoded = d.decrypt(e.encrypt(CLEARTEXT))
3038
assert decoded == CLEARTEXT
39+
40+
def test_encrypt_message_attr(self):
41+
"""Test that the Message object has accessible Python attributes (mac, ciphertext, ephemeral_key)."""
42+
decryption = PkDecryption()
43+
encryption = PkEncryption.from_key(decryption.public_key)
44+
45+
message = encryption.encrypt(CLEARTEXT)
46+
47+
assert message.mac is not None
48+
assert message.ciphertext is not None
49+
assert message.ephemeral_key is not None
50+
51+
52+
def test_message_from_invalid_base64(self):
53+
"""Test that invalid base64 input raises PkDecodeException."""
54+
# Test invalid ciphertext
55+
with pytest.raises(PkDecodeException, match="Invalid symbol"):
56+
Message.from_base64(
57+
"not-valid-base64!@#", # Invalid base64 for ciphertext
58+
base64.b64encode(b"some_mac").decode(), # Valid base64
59+
base64.b64encode(b"some_key").decode() # Valid base64
60+
)
61+
62+
# Test invalid mac
63+
with pytest.raises(PkDecodeException, match="Invalid symbol"):
64+
Message.from_base64(
65+
base64.b64encode(b"some_text").decode(),
66+
"not-valid-base64!@#", # Invalid base64 for mac
67+
base64.b64encode(b"some_key").decode()
68+
)
69+
70+
# Test invalid ephemeral key
71+
with pytest.raises(PkDecodeException, match="Invalid symbol"):
72+
Message.from_base64(
73+
base64.b64encode(b"some_text").decode(),
74+
base64.b64encode(b"some_mac").decode(),
75+
"not-valid-base64!@#" # Invalid base64 for ephemeral key
76+
)

0 commit comments

Comments
 (0)