libsoliton/soliton_py/src/stream.rs

//! 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(())
}