synta-python 0.2.5

//! Python binding for the CMS EnvelopedData fluent builder:
//! [`PyEnvelopedDataBuilder`] (RFC 5652 §6).

use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyList};

use super::enveloped::PyEnvelopedData;

// ── PyEnvelopedDataBuilder ────────────────────────────────────────────────────

/// Builder for CMS ``EnvelopedData`` (RFC 5652 §6).
///
/// Accumulates the cryptographic ingredients during construction and optional
/// ``OriginatorInfo`` / ``UnprotectedAttributes`` fields added afterwards, then
/// assembles the final DER on :meth:`build`.
///
/// ```python,ignore
/// import synta, synta.cms
///
/// # Simple: one recipient, AES-256-CBC content encryption (default).
/// ed = (
///     synta.cms.EnvelopedDataBuilder(plaintext, [(cert, synta.RSAES_OAEP)])
///     .build()
/// )
///
/// # With an originator certificate:
/// ed = (
///     synta.cms.EnvelopedDataBuilder(plaintext, [(cert, synta.RSAES_OAEP)])
///     .add_originator_cert(signer_cert)   # Certificate object or bytes
///     .build()
/// )
/// ```
///
/// Each setter returns the **same** builder object, enabling chaining.
#[pyclass(name = "EnvelopedDataBuilder")]
pub struct PyEnvelopedDataBuilder {
    /// Plaintext to encrypt.
    plaintext: Vec<u8>,
    /// Content-encryption algorithm OID components.
    content_enc_alg_comps: Vec<u32>,
    /// Recipients: (cert_der, key_wrap_oid_components).
    recipients: Vec<(Vec<u8>, Vec<u32>)>,
    /// Originator certificates: each is a full DER Certificate SEQUENCE.
    originator_certs: Vec<Vec<u8>>,
    /// Originator CRLs: each is a full DER CertificateList SEQUENCE.
    originator_crls: Vec<Vec<u8>>,
    /// Raw content bytes for `UnprotectedAttributes` (without outer SET tag/len).
    unprotected_attrs: Option<Vec<u8>>,
}

#[pymethods]
impl PyEnvelopedDataBuilder {
    /// Create an ``EnvelopedDataBuilder``.
    ///
    /// ``plaintext`` is the content to encrypt.
    ///
    /// ``recipients`` is a non-empty list of ``(cert, key_wrap_oid)`` tuples.
    /// Each ``cert`` may be a :class:`~synta.Certificate` object or a
    /// ``bytes`` value containing a DER-encoded ``Certificate`` SEQUENCE.
    /// ``key_wrap_oid`` is the :class:`~synta.ObjectIdentifier` (or
    /// equivalent) that selects the key-transport algorithm; the supported
    /// values are ``synta.cms.ID_RSAES_OAEP`` (RSA-OAEP / SHA-256, recommended)
    /// and ``synta.cms.ID_RSA_ENCRYPTION`` (RSA PKCS#1 v1.5, legacy).
    ///
    /// ``content_enc_alg`` (keyword-only) selects the content-encryption
    /// algorithm OID; defaults to ``id-aes256-CBC``.
    ///
    /// Raises :exc:`ValueError` if ``recipients`` is empty.
    /// Raises :exc:`TypeError` if any element of ``recipients`` is not a
    /// two-element tuple or if ``cert`` is not a ``Certificate`` or ``bytes``.
    #[new]
    #[pyo3(signature = (plaintext, recipients, *, content_enc_alg = None))]
    fn new(
        py: Python<'_>,
        plaintext: &[u8],
        recipients: &Bound<'_, PyList>,
        content_enc_alg: Option<&Bound<'_, PyAny>>,
    ) -> PyResult<Self> {
        use pyo3::types::PyTuple;

        let content_enc_alg_comps: Vec<u32> = match content_enc_alg {
            Some(obj) => super::super::oid_from_pyany(obj)?.components().to_vec(),
            None => synta_certificate::pkcs12_types::ID_AES256_CBC.to_vec(),
        };

        let mut recipient_pairs: Vec<(Vec<u8>, Vec<u32>)> = Vec::with_capacity(recipients.len());
        for item in recipients.iter() {
            let tup = item.cast::<PyTuple>().map_err(|_| {
                pyo3::exceptions::PyTypeError::new_err(
                    "each recipient entry must be a (cert, key_wrap_oid) tuple",
                )
            })?;
            if tup.len() != 2 {
                return Err(pyo3::exceptions::PyTypeError::new_err(
                    "each recipient tuple must have exactly 2 elements: (cert, key_wrap_oid)",
                ));
            }
            let cert_item = tup.get_item(0)?;
            let cert_der: Vec<u8> =
                if let Ok(py_cert) = cert_item.cast::<super::super::cert::PyCertificate>() {
                    py_cert.get().raw.to_vec()
                } else if let Ok(py_bytes) = cert_item.cast::<PyBytes>() {
                    py_bytes.as_bytes().to_vec()
                } else {
                    return Err(pyo3::exceptions::PyTypeError::new_err(
                        "cert must be a synta.Certificate or bytes",
                    ));
                };
            let kw_comps = super::super::oid_from_pyany(&tup.get_item(1)?)?
                .components()
                .to_vec();
            recipient_pairs.push((cert_der, kw_comps));
        }
        if recipient_pairs.is_empty() {
            return Err(pyo3::exceptions::PyValueError::new_err(
                "recipients list must not be empty",
            ));
        }
        let _ = py;
        Ok(Self {
            plaintext: plaintext.to_vec(),
            content_enc_alg_comps,
            recipients: recipient_pairs,
            originator_certs: Vec::new(),
            originator_crls: Vec::new(),
            unprotected_attrs: None,
        })
    }

    /// Add an originator certificate to ``OriginatorInfo.certs``.
    ///
    /// ``cert`` may be a :class:`~synta.Certificate` object or raw DER
    /// ``bytes`` of a ``Certificate`` SEQUENCE.
    fn add_originator_cert<'py>(
        slf: Bound<'py, Self>,
        cert: &Bound<'_, PyAny>,
    ) -> PyResult<Bound<'py, Self>> {
        let cert_der: Vec<u8> =
            if let Ok(py_cert) = cert.cast::<super::super::cert::PyCertificate>() {
                py_cert.get().raw.to_vec()
            } else if let Ok(py_bytes) = cert.cast::<PyBytes>() {
                py_bytes.as_bytes().to_vec()
            } else {
                return Err(pyo3::exceptions::PyTypeError::new_err(
                    "cert must be a synta.Certificate or bytes",
                ));
            };
        slf.borrow_mut().originator_certs.push(cert_der);
        Ok(slf)
    }

    /// Add a CRL to ``OriginatorInfo.crls``.
    ///
    /// ``crl`` may be a :class:`~synta.CertificateList` object or raw DER
    /// ``bytes`` of a ``CertificateList`` SEQUENCE.
    fn add_originator_crl<'py>(
        slf: Bound<'py, Self>,
        crl: &Bound<'_, PyAny>,
    ) -> PyResult<Bound<'py, Self>> {
        let crl_der: Vec<u8> = if let Ok(py_crl) = crl.cast::<super::super::pkix::PyCrl>() {
            py_crl.get().raw.to_vec()
        } else if let Ok(py_bytes) = crl.cast::<PyBytes>() {
            py_bytes.as_bytes().to_vec()
        } else {
            return Err(pyo3::exceptions::PyTypeError::new_err(
                "crl must be a synta.CertificateList or bytes",
            ));
        };
        slf.borrow_mut().originator_crls.push(crl_der);
        Ok(slf)
    }

    /// Set the ``unprotectedAttrs [1] IMPLICIT`` field.
    ///
    /// ``attrs`` must be the raw **content** bytes of the ``SET OF Attribute``
    /// (without the outer SET tag or length prefix).
    ///
    /// Returns the same builder object to enable chaining.
    fn set_unprotected_attrs<'py>(slf: Bound<'py, Self>, attrs: &[u8]) -> Bound<'py, Self> {
        slf.borrow_mut().unprotected_attrs = Some(attrs.to_vec());
        slf
    }

    /// Assemble and return the DER-encoded ``EnvelopedData`` as a
    /// :class:`~synta.cms.EnvelopedData` object.
    ///
    /// Generates a fresh content-encryption key (CEK), encrypts the
    /// plaintext, wraps the CEK independently for each recipient, then
    /// assembles the ``EnvelopedData`` SEQUENCE from all pre-computed
    /// ingredients.
    ///
    /// Raises :exc:`ValueError` on unsupported key-wrap OIDs or crypto
    /// backend errors.
    fn build(&self, py: Python<'_>) -> PyResult<PyEnvelopedData> {
        use synta_certificate::KeyWrapAlgorithm;

        let oid_for = |comps: &[u32]| -> PyResult<KeyWrapAlgorithm> {
            if comps == synta_certificate::oids::RSAES_OAEP {
                Ok(KeyWrapAlgorithm::RsaOaepSha256)
            } else if comps == synta_certificate::oids::RSA_ENCRYPTION {
                Ok(KeyWrapAlgorithm::RsaPkcs1v15)
            } else {
                Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "unsupported key-wrap OID {:?}; use RSAES_OAEP or RSA_ENCRYPTION",
                    comps,
                )))
            }
        };

        // Map stored (cert_der, kw_comps) into (cert_der_slice, KeyWrapAlgorithm).
        let kw_pairs: Vec<KeyWrapAlgorithm> = self
            .recipients
            .iter()
            .map(|(_, kw_comps)| oid_for(kw_comps))
            .collect::<PyResult<_>>()?;
        let pairs_ref: Vec<(&[u8], KeyWrapAlgorithm)> = self
            .recipients
            .iter()
            .zip(kw_pairs.iter())
            .map(|((cert_der, _), kw)| (cert_der.as_slice(), *kw))
            .collect();

        // 1. Generate CEK, encrypt plaintext, build KTRIs → pre-loaded builder.
        let mut builder = synta_certificate::default_prepare_enveloped_data(
            &self.plaintext,
            &pairs_ref,
            &self.content_enc_alg_comps,
        )
        .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;

        // 2. Attach optional OriginatorInfo and UnprotectedAttributes.
        for cert_der in &self.originator_certs {
            builder = builder.originator_cert(cert_der);
        }
        for crl_der in &self.originator_crls {
            builder = builder.originator_crl(crl_der);
        }
        if let Some(attrs) = &self.unprotected_attrs {
            builder = builder.unprotected_attrs(attrs);
        }

        // 3. Assemble and return.
        let der = builder
            .build()
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
        let py_bytes = PyBytes::new(py, &der).unbind();
        PyEnvelopedData::from_der(py, py_bytes.into_bound(py))
    }
}