initial commit

Signed-off-by: Kamal Tufekcic <kamal@lo.sh>

soliton_py/src/stream.rs

@@ -0,0 +1,241 @@
+//! Streaming AEAD: chunked encryption/decryption for files and media.
+
+use pyo3::prelude::*;
+use pyo3::types::PyBytes;
+
+use crate::errors::to_py_err;
+
+/// Streaming encryptor for chunked file/media encryption.
+///
+/// Produces a 26-byte header (sent first) and one ciphertext chunk per call
+/// to ``encrypt_chunk``. Non-final chunks must be exactly 1 MiB of plaintext;
+/// the final chunk may be shorter.
+///
+/// Use as a context manager::
+///
+///     with soliton.StreamEncryptor(key) as enc:
+///         header = enc.header()
+///         ct1 = enc.encrypt_chunk(data1)
+///         ct2 = enc.encrypt_chunk(data2, is_last=True)
+#[pyclass]
+pub struct StreamEncryptor {
+    inner: Option<soliton::streaming::StreamEncryptor>,
+}
+
+#[pymethods]
+impl StreamEncryptor {
+    /// Create a streaming encryptor.
+    ///
+    /// Args:
+    ///     key: 32-byte encryption key.
+    ///     aad: Optional additional authenticated data (bound to all chunks).
+    ///     compress: Enable zstd compression per chunk (default False).
+    #[new]
+    #[pyo3(signature = (key, aad=None, compress=false))]
+    fn new(key: &[u8], aad: Option<&[u8]>, compress: bool) -> PyResult<Self> {
+        if key.len() != 32 {
+            return Err(crate::errors::InvalidLengthError::new_err(
+                "key must be 32 bytes",
+            ));
+        }
+        let key_arr: &[u8; 32] = key.try_into().unwrap();
+        let aad_bytes = aad.unwrap_or(&[]);
+        let enc = soliton::streaming::stream_encrypt_init(key_arr, aad_bytes, compress)
+            .map_err(to_py_err)?;
+        Ok(Self { inner: Some(enc) })
+    }
+
+    /// The 26-byte stream header. Send this before any chunks.
+    fn header<'py>(&self, py: Python<'py>) -> PyResult<Py<PyBytes>> {
+        let inner = self
+            .inner
+            .as_ref()
+            .ok_or_else(|| crate::errors::InvalidDataError::new_err("encryptor closed"))?;
+        Ok(PyBytes::new(py, &inner.header()).into())
+    }
+
+    /// Encrypt one chunk.
+    ///
+    /// Args:
+    ///     plaintext: Chunk data. Non-final chunks must be exactly 1,048,576 bytes.
+    ///     is_last: True for the final chunk.
+    ///
+    /// Returns:
+    ///     Encrypted chunk bytes.
+    #[pyo3(signature = (plaintext, is_last=false))]
+    fn encrypt_chunk<'py>(
+        &mut self,
+        py: Python<'py>,
+        plaintext: &[u8],
+        is_last: bool,
+    ) -> PyResult<Py<PyBytes>> {
+        let inner = self
+            .inner
+            .as_mut()
+            .ok_or_else(|| crate::errors::InvalidDataError::new_err("encryptor closed"))?;
+        let ct = inner.encrypt_chunk(plaintext, is_last).map_err(to_py_err)?;
+        Ok(PyBytes::new(py, &ct).into())
+    }
+
+    /// Encrypt a specific chunk by index (random access, stateless).
+    fn encrypt_chunk_at<'py>(
+        &self,
+        py: Python<'py>,
+        index: u64,
+        plaintext: &[u8],
+        is_last: bool,
+    ) -> PyResult<Py<PyBytes>> {
+        let inner = self
+            .inner
+            .as_ref()
+            .ok_or_else(|| crate::errors::InvalidDataError::new_err("encryptor closed"))?;
+        let ct = inner
+            .encrypt_chunk_at(index, is_last, plaintext)
+            .map_err(to_py_err)?;
+        Ok(PyBytes::new(py, &ct).into())
+    }
+
+    /// Whether the encryptor has been finalized.
+    fn is_finalized(&self) -> PyResult<bool> {
+        let inner = self
+            .inner
+            .as_ref()
+            .ok_or_else(|| crate::errors::InvalidDataError::new_err("encryptor closed"))?;
+        Ok(inner.is_finalized())
+    }
+
+    fn close(&mut self) {
+        self.inner = None;
+    }
+
+    fn __enter__(slf: Py<Self>) -> Py<Self> {
+        slf
+    }
+
+    fn __exit__(
+        &mut self,
+        _exc_type: Option<&Bound<'_, PyAny>>,
+        _exc_val: Option<&Bound<'_, PyAny>>,
+        _exc_tb: Option<&Bound<'_, PyAny>>,
+    ) {
+        self.close();
+    }
+}
+
+/// Streaming decryptor for chunked file/media decryption.
+///
+/// Use as a context manager::
+///
+///     with soliton.StreamDecryptor(key, header) as dec:
+///         pt1, is_last = dec.decrypt_chunk(ct1)
+///         pt2, is_last = dec.decrypt_chunk(ct2)
+#[pyclass]
+pub struct StreamDecryptor {
+    inner: Option<soliton::streaming::StreamDecryptor>,
+}
+
+#[pymethods]
+impl StreamDecryptor {
+    /// Create a streaming decryptor.
+    ///
+    /// Args:
+    ///     key: 32-byte decryption key (same key used for encryption).
+    ///     header: 26-byte stream header from the encryptor.
+    ///     aad: Optional additional authenticated data (must match encryption).
+    #[new]
+    #[pyo3(signature = (key, header, aad=None))]
+    fn new(key: &[u8], header: &[u8], aad: Option<&[u8]>) -> PyResult<Self> {
+        if key.len() != 32 {
+            return Err(crate::errors::InvalidLengthError::new_err(
+                "key must be 32 bytes",
+            ));
+        }
+        if header.len() != 26 {
+            return Err(crate::errors::InvalidLengthError::new_err(
+                "header must be 26 bytes",
+            ));
+        }
+        let key_arr: &[u8; 32] = key.try_into().unwrap();
+        let hdr_arr: &[u8; 26] = header.try_into().unwrap();
+        let aad_bytes = aad.unwrap_or(&[]);
+        let dec = soliton::streaming::stream_decrypt_init(key_arr, hdr_arr, aad_bytes)
+            .map_err(to_py_err)?;
+        Ok(Self { inner: Some(dec) })
+    }
+
+    /// Decrypt the next sequential chunk.
+    ///
+    /// Returns:
+    ///     Tuple of (plaintext: bytes, is_last: bool).
+    fn decrypt_chunk<'py>(
+        &mut self,
+        py: Python<'py>,
+        chunk: &[u8],
+    ) -> PyResult<(Py<PyBytes>, bool)> {
+        let inner = self
+            .inner
+            .as_mut()
+            .ok_or_else(|| crate::errors::InvalidDataError::new_err("decryptor closed"))?;
+        let (pt, is_last) = inner.decrypt_chunk(chunk).map_err(to_py_err)?;
+        Ok((PyBytes::new(py, &pt).into(), is_last))
+    }
+
+    /// Decrypt a specific chunk by index (random access).
+    ///
+    /// Returns:
+    ///     Tuple of (plaintext: bytes, is_last: bool).
+    fn decrypt_chunk_at<'py>(
+        &self,
+        py: Python<'py>,
+        index: u64,
+        chunk: &[u8],
+    ) -> PyResult<(Py<PyBytes>, bool)> {
+        let inner = self
+            .inner
+            .as_ref()
+            .ok_or_else(|| crate::errors::InvalidDataError::new_err("decryptor closed"))?;
+        let (pt, is_last) = inner.decrypt_chunk_at(index, chunk).map_err(to_py_err)?;
+        Ok((PyBytes::new(py, &pt).into(), is_last))
+    }
+
+    /// Whether the decryptor has seen the final chunk.
+    fn is_finalized(&self) -> PyResult<bool> {
+        let inner = self
+            .inner
+            .as_ref()
+            .ok_or_else(|| crate::errors::InvalidDataError::new_err("decryptor closed"))?;
+        Ok(inner.is_finalized())
+    }
+
+    /// Next expected sequential chunk index.
+    fn expected_index(&self) -> PyResult<u64> {
+        let inner = self
+            .inner
+            .as_ref()
+            .ok_or_else(|| crate::errors::InvalidDataError::new_err("decryptor closed"))?;
+        Ok(inner.expected_index())
+    }
+
+    fn close(&mut self) {
+        self.inner = None;
+    }
+
+    fn __enter__(slf: Py<Self>) -> Py<Self> {
+        slf
+    }
+
+    fn __exit__(
+        &mut self,
+        _exc_type: Option<&Bound<'_, PyAny>>,
+        _exc_val: Option<&Bound<'_, PyAny>>,
+        _exc_tb: Option<&Bound<'_, PyAny>>,
+    ) {
+        self.close();
+    }
+}
+
+pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<StreamEncryptor>()?;
+    m.add_class::<StreamDecryptor>()?;
+    Ok(())
+}