diff --git a/Cargo.toml b/Cargo.toml index 58668b8..8533369 100644 --- a/Cargo.toml +++ b/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", rev = "12f9036bf7f2536c172273602afcdc9aeddf8cf7" } +vodozemac = { git = "https://github.com/matrix-org/vodozemac.git", rev = "12f9036bf7f2536c172273602afcdc9aeddf8cf7", features = ["insecure-pk-encryption"] } [package.metadata.maturin] name = "vodozemac" diff --git a/src/error.rs b/src/error.rs index 0147656..728c0a5 100644 --- a/src/error.rs +++ b/src/error.rs @@ -139,3 +139,29 @@ impl From 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 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()), + } + } +} diff --git a/src/lib.rs b/src/lib.rs index a0084e8..556d01e 100644 --- a/src/lib.rs +++ b/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::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; + + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; m.add("KeyException", py.get_type_bound::())?; m.add( @@ -55,6 +61,14 @@ fn my_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { "MegolmDecryptionException", py.get_type_bound::(), )?; + m.add( + "PkInvalidKeySizeException", + py.get_type_bound::(), + )?; + m.add( + "PkDecodeException", + py.get_type_bound::(), + )?; Ok(()) } diff --git a/src/pk_encryption.rs b/src/pk_encryption.rs new file mode 100644 index 0000000..fcc7480 --- /dev/null +++ b/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, + /// The message authentication code of the message. + /// + /// *Warning*: As stated in the module description, this does not + /// authenticate the message. + mac: Vec, + /// The ephemeral Curve25519PublicKey of the message which was used to + /// derive the individual message key. + ephemeral_key: Vec, +} + +/// ☣️ 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 { + 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, 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 { + 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(), + } + } +} diff --git a/src/types/curve25519.rs b/src/types/curve25519.rs index cbfc621..11ff81f 100644 --- a/src/types/curve25519.rs +++ b/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,10 +26,29 @@ impl Curve25519PublicKey { }) } + #[classmethod] + pub fn from_bytes(_cls: &Bound<'_, PyType>, bytes: &[u8]) -> Result { + 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 { + convert_to_pybytes(self.inner.to_bytes()) + } + #[classattr] const __hash__: Option = None; @@ -33,3 +56,71 @@ impl Curve25519PublicKey { self.inner == other.inner } } + +#[pyclass] +#[derive(Clone)] +pub struct Curve25519SecretKey { + pub(crate) inner: vodozemac::Curve25519SecretKey, +} + +impl From 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::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 { + 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 { + 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 = None; + + fn __eq__(&self, other: &Self) -> bool { + self.inner.to_bytes() == other.inner.to_bytes() + } +} diff --git a/tests/pk_encryption_test.py b/tests/pk_encryption_test.py new file mode 100644 index 0000000..2ce2195 --- /dev/null +++ b/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 \ No newline at end of file