synta-python 0.1.9

//! Python binding for the CMS SignedData fluent builder:
//! [`PySignedDataBuilder`] (RFC 5652 §5).

use std::str::FromStr;

use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use synta::types::string::OctetStringRef;
use synta::{Element, ExplicitTag, Integer, Null, ObjectIdentifier, RawDer, SetOf, Tag};
use synta::{Encode, Encoder, Encoding};
use synta_certificate::cms_rfc5652_types::{
    Attribute, EncapsulatedContentInfo, SignedData, SignerInfo,
};
use synta_certificate::pkcs7_types::ContentInfo;
use synta_certificate::AlgorithmIdentifier;
use synta_certificate::{DataHasher as _, PrivateKey as _};

use crate::crypto_keys::PyPrivateKey;

// ── digest algorithm helpers ──────────────────────────────────────────────────

/// Return the OID string for a named hash algorithm, or `None` for unknown names.
fn digest_alg_oid(name: &str) -> Option<&'static str> {
    match name {
        "sha1" => Some("1.3.14.3.2.26"),
        "sha256" => Some("2.16.840.1.101.3.4.2.1"),
        "sha384" => Some("2.16.840.1.101.3.4.2.2"),
        "sha512" => Some("2.16.840.1.101.3.4.2.3"),
        _ => None,
    }
}

/// Build an [`AlgorithmIdentifier`] for a named hash algorithm with `NULL` parameters.
fn make_digest_alg(name: &str) -> Option<AlgorithmIdentifier<'static>> {
    let oid_str = digest_alg_oid(name)?;
    let oid = ObjectIdentifier::from_str(oid_str).ok()?;
    Some(AlgorithmIdentifier {
        algorithm: oid,
        parameters: Some(Element::Null(Null)),
    })
}

// ── IssuerAndSerialNumber encoding ────────────────────────────────────────────

/// Parse `cert_der` and return the DER-encoded `IssuerAndSerialNumber` SEQUENCE.
fn build_iasn_der(cert_der: &[u8]) -> PyResult<Vec<u8>> {
    use synta_certificate::cms_2010_types::IssuerAndSerialNumber;
    use synta_certificate::{Certificate, Name};

    let cert: Certificate = synta::Decoder::new(cert_der, Encoding::Ber)
        .decode()
        .map_err(|e| PyValueError::new_err(format!("failed to parse signing certificate: {e}")))?;

    let issuer_raw = cert.tbs_certificate.issuer.as_bytes();
    let issuer = Name::from_der(issuer_raw)
        .map_err(|e| PyValueError::new_err(format!("issuer decode: {e}")))?;
    let serial_number = Integer::from_bytes(cert.tbs_certificate.serial_number.as_bytes());

    IssuerAndSerialNumber {
        issuer,
        serial_number,
    }
    .to_der()
    .map_err(|e| PyValueError::new_err(format!("IssuerAndSerialNumber encode: {e}")))
}

// ── signed attributes ─────────────────────────────────────────────────────────

/// Encode a SET OF containing a single encodable value, return the SET TLV bytes.
fn encode_set_of_one<T: Encode>(value: T) -> PyResult<Vec<u8>> {
    let mut set: SetOf<T> = SetOf::new();
    set.push(value);
    let mut enc = Encoder::new(Encoding::Der);
    set.encode(&mut enc)
        .map_err(|e| PyValueError::new_err(format!("{e}")))?;
    enc.finish()
        .map_err(|e| PyValueError::new_err(format!("{e}")))
}

/// Build the DER bytes for the `contentType` signed attribute (OID 1.2.840.113549.1.9.3).
fn build_content_type_attr(content_oid: &ObjectIdentifier) -> PyResult<Vec<u8>> {
    let attr_type =
        ObjectIdentifier::from_str("1.2.840.113549.1.9.3").expect("contentType OID is valid");
    let set_bytes = encode_set_of_one(content_oid.clone())?;
    Attribute {
        attr_type,
        attr_values: RawDer(&set_bytes),
    }
    .to_der()
    .map_err(|e| PyValueError::new_err(format!("{e}")))
}

/// Build the DER bytes for the `messageDigest` signed attribute (OID 1.2.840.113549.1.9.4).
fn build_message_digest_attr(digest: &[u8]) -> PyResult<Vec<u8>> {
    let attr_type =
        ObjectIdentifier::from_str("1.2.840.113549.1.9.4").expect("messageDigest OID is valid");
    let set_bytes = encode_set_of_one(OctetStringRef::new(digest))?;
    Attribute {
        attr_type,
        attr_values: RawDer(&set_bytes),
    }
    .to_der()
    .map_err(|e| PyValueError::new_err(format!("{e}")))
}

// ── SignerSpec ────────────────────────────────────────────────────────────────

/// Internal per-signer configuration: key, certificate, and digest algorithm.
struct SignerSpec {
    key: synta_certificate::BackendPrivateKey,
    cert_der: Vec<u8>,
    hash_algorithm: String,
}

// ── PySignedDataBuilder ───────────────────────────────────────────────────────

/// Builder for CMS ``SignedData`` (RFC 5652 §5).
///
/// Accumulates the signing ingredients and optional extra certificates, then
/// assembles the final DER on :meth:`build`.
///
/// ```python,ignore
/// import synta, synta.cms
///
/// # Attached (content included in SignedData):
/// sd_der = (
///     synta.cms.SignedDataBuilder(content_bytes, content_type="1.2.840.113549.1.7.1")
///     .add_signer(private_key, signing_cert_der, hash_algorithm="sha256")
///     .add_cert(intermediate_cert_der)
///     .build()
/// )
///
/// # Detached (content excluded from SignedData):
/// sd_der = (
///     synta.cms.SignedDataBuilder(content_bytes, detached=True)
///     .add_signer(private_key, signing_cert_der, hash_algorithm="sha256")
///     .build()
/// )
/// ```
///
/// Each ``add_signer`` / ``add_cert`` call returns the **same** builder object,
/// enabling chaining.
#[pyclass(name = "SignedDataBuilder")]
pub struct PySignedDataBuilder {
    /// Content to be signed (the ``eContent`` bytes).
    content: Vec<u8>,
    /// ``eContentType`` OID for ``EncapsulatedContentInfo``.
    content_oid: ObjectIdentifier,
    /// When ``True`` the ``eContent`` field is omitted from the output.
    detached: bool,
    /// Per-signer specifications.
    signers: Vec<SignerSpec>,
    /// Extra certificates to include in the ``certificates [0]`` field.
    extra_certs: Vec<Vec<u8>>,
}

#[pymethods]
impl PySignedDataBuilder {
    /// Create a ``SignedDataBuilder``.
    ///
    /// ``content`` is the data that will be signed.
    ///
    /// ``content_type`` (keyword-only) is the dotted-decimal OID string for
    /// ``eContentType``; defaults to ``id-data`` (``1.2.840.113549.1.7.1``).
    ///
    /// ``detached`` (keyword-only) controls whether the content is embedded in
    /// the output (``False``, default) or omitted (``True`` — detached signature).
    ///
    /// Raises :exc:`ValueError` if ``content_type`` is not a valid OID string.
    #[new]
    #[pyo3(signature = (content, *, content_type = None, detached = false))]
    fn new(content: &[u8], content_type: Option<&str>, detached: bool) -> PyResult<Self> {
        let oid_str = content_type.unwrap_or("1.2.840.113549.1.7.1");
        let content_oid = ObjectIdentifier::from_str(oid_str).map_err(|e| {
            PyValueError::new_err(format!("invalid content_type OID {oid_str:?}: {e}"))
        })?;
        Ok(Self {
            content: content.to_vec(),
            content_oid,
            detached,
            signers: Vec::new(),
            extra_certs: Vec::new(),
        })
    }

    /// Add a signer to the ``SignedData``.
    ///
    /// ``key`` is the private key used for signing.
    ///
    /// ``cert_der`` is the DER-encoded signing certificate; the issuer name and
    /// serial number are extracted from it to populate ``IssuerAndSerialNumber``,
    /// and the certificate itself is included in the ``certificates`` field.
    ///
    /// ``hash_algorithm`` (keyword-only) is the digest algorithm name; must be
    /// one of ``"sha1"``, ``"sha256"`` (default), ``"sha384"``, or ``"sha512"``.
    ///
    /// Returns the same builder to allow chaining.
    ///
    /// Raises :exc:`ValueError` for unsupported hash algorithms.
    #[pyo3(signature = (key, cert_der, hash_algorithm = "sha256"))]
    fn add_signer<'py>(
        slf: Bound<'py, Self>,
        key: &PyPrivateKey,
        cert_der: &[u8],
        hash_algorithm: &str,
    ) -> PyResult<Bound<'py, Self>> {
        match hash_algorithm {
            "sha1" | "sha256" | "sha384" | "sha512" => {}
            other => {
                return Err(PyValueError::new_err(format!(
                    "unsupported hash_algorithm {other:?}; use sha1, sha256, sha384, or sha512"
                )));
            }
        }
        let spec = SignerSpec {
            key: key.inner.clone(),
            cert_der: cert_der.to_vec(),
            hash_algorithm: hash_algorithm.to_owned(),
        };
        slf.borrow_mut().signers.push(spec);
        Ok(slf)
    }

    /// Add an extra certificate to the ``certificates [0]`` field.
    ///
    /// ``cert_der`` is the DER-encoded certificate to include.
    ///
    /// Returns the same builder to allow chaining.
    fn add_cert<'py>(slf: Bound<'py, Self>, cert_der: &[u8]) -> Bound<'py, Self> {
        slf.borrow_mut().extra_certs.push(cert_der.to_vec());
        slf
    }

    /// Assemble and return the DER-encoded ``ContentInfo`` wrapping a
    /// ``SignedData`` as :class:`bytes`.
    ///
    /// For each signer:
    ///
    /// 1. Computes the message digest of the content.
    /// 2. Builds and DER-encodes the ``signedAttrs`` (``contentType`` +
    ///    ``messageDigest``).
    /// 3. Signs the ``signedAttrs`` bytes.
    /// 4. Assembles the ``SignerInfo`` SEQUENCE via the generated ASN.1 type.
    ///
    /// Then constructs ``SignedData`` and wraps it in a ``ContentInfo``.
    ///
    /// Raises :exc:`ValueError` on certificate parse errors, unsupported
    /// algorithms, or crypto backend failures.
    fn build<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        if self.signers.is_empty() {
            return Err(PyValueError::new_err(
                "at least one signer must be added before calling build()",
            ));
        }

        let hasher = synta_certificate::default_data_hasher();

        // ── Phase 1: per-signer computation ──────────────────────────────────
        // Collect intermediate owned buffers for each signer.  These stay alive
        // for the full duration of `build`, so we can safely borrow from them
        // when constructing the generated structs in Phase 2.
        struct SignerBuffers {
            iasn_der: Vec<u8>,
            signed_attrs_content: Vec<u8>,
            sig_alg_der: Vec<u8>,
            signature: Vec<u8>,
            cert_der: Vec<u8>,
            hash_algorithm: String,
        }

        let mut signer_bufs: Vec<SignerBuffers> = Vec::with_capacity(self.signers.len());

        for spec in &self.signers {
            // IssuerAndSerialNumber SEQUENCE bytes
            let iasn_der = build_iasn_der(&spec.cert_der)?;

            // Digest of the content
            let digest_bytes = hasher
                .hash_data(&spec.hash_algorithm, &self.content)
                .map_err(|e| PyValueError::new_err(format!("hash failed: {e}")))?;

            // signedAttrs SET content (attribute SEQUENCE bytes concatenated)
            // RFC 5652 §5.4: the signed content is the SET-tagged form (0x31)
            let content_type_attr = build_content_type_attr(&self.content_oid)?;
            let msg_digest_attr = build_message_digest_attr(&digest_bytes)?;
            let signed_attrs_content: Vec<u8> =
                [content_type_attr.as_slice(), msg_digest_attr.as_slice()].concat();

            // For signing: re-tag as SET (0x31) per RFC 5652 §5.4
            let to_sign = {
                let mut enc = Encoder::new(Encoding::Der);
                enc.write_tag(Tag::universal_constructed(17))
                    .map_err(|e| PyValueError::new_err(format!("{e}")))?;
                enc.write_length(signed_attrs_content.len())
                    .map_err(|e| PyValueError::new_err(format!("{e}")))?;
                enc.write_bytes(&signed_attrs_content);
                enc.finish()
                    .map_err(|e| PyValueError::new_err(format!("{e}")))?
            };

            // Sign the SET-tagged signedAttrs
            let signer = spec.key.as_signer(&spec.hash_algorithm);
            let sig_alg_der = signer
                .signature_algorithm_der_erased()
                .map_err(|e| PyValueError::new_err(format!("sig alg: {e}")))?;
            let signature = signer
                .sign_tbs_erased(&to_sign)
                .map_err(|e| PyValueError::new_err(format!("sign: {e}")))?;

            signer_bufs.push(SignerBuffers {
                iasn_der,
                signed_attrs_content,
                sig_alg_der,
                signature,
                cert_der: spec.cert_der.clone(),
                hash_algorithm: spec.hash_algorithm.clone(),
            });
        }

        // ── Phase 2: assemble SignerInfo structs ──────────────────────────────
        // Encode each SignerInfo to DER immediately so we can collect the bytes.
        let mut si_ders: Vec<Vec<u8>> = Vec::with_capacity(signer_bufs.len());
        let mut seen_hash_algs: Vec<String> = Vec::new();
        let mut digest_algs: SetOf<AlgorithmIdentifier<'static>> = SetOf::new();

        for buf in &signer_bufs {
            let digest_alg = make_digest_alg(&buf.hash_algorithm).ok_or_else(|| {
                PyValueError::new_err(format!(
                    "unsupported hash_algorithm {:?}",
                    buf.hash_algorithm
                ))
            })?;

            if !seen_hash_algs.contains(&buf.hash_algorithm) {
                seen_hash_algs.push(buf.hash_algorithm.clone());
                digest_algs.push(digest_alg.clone());
            }

            // Parse the signature AlgorithmIdentifier from its DER encoding.
            // This borrows from `buf.sig_alg_der` which lives for the duration
            // of the function.
            let sig_alg: AlgorithmIdentifier<'_> =
                synta::Decoder::new(&buf.sig_alg_der, Encoding::Der)
                    .decode()
                    .map_err(|e| PyValueError::new_err(format!("sig alg decode: {e}")))?;

            // `signed_attrs` RawDer holds only the SET VALUE bytes (no outer tag).
            // The Asn1Sequence derive for SignerInfo writes `a0 <len> <bytes>` around them.
            let si = SignerInfo {
                version: Integer::from(1),
                sid: RawDer(&buf.iasn_der),
                digest_algorithm: digest_alg,
                signed_attrs: Some(RawDer(&buf.signed_attrs_content)),
                signature_algorithm: sig_alg,
                signature: OctetStringRef::new(&buf.signature),
                unsigned_attrs: None,
            };

            let si_der = si
                .to_der()
                .map_err(|e| PyValueError::new_err(format!("SignerInfo encode: {e}")))?;
            si_ders.push(si_der);
        }

        // ── Phase 3: assemble SignedData ──────────────────────────────────────

        // Signer infos: re-parse from DER so the SignedData struct can own them.
        let mut signer_infos: SetOf<SignerInfo<'_>> = SetOf::new();
        for si_der in &si_ders {
            let si: SignerInfo<'_> = SignerInfo::from_der(si_der)
                .map_err(|e| PyValueError::new_err(format!("SignerInfo re-parse: {e}")))?;
            signer_infos.push(si);
        }

        // EncapsulatedContentInfo
        let eci = if self.detached {
            EncapsulatedContentInfo {
                e_content_type: self.content_oid.clone(),
                e_content: None,
            }
        } else {
            EncapsulatedContentInfo {
                e_content_type: self.content_oid.clone(),
                e_content: Some(OctetStringRef::new(&self.content)),
            }
        };

        // certificates [0] IMPLICIT — concatenated cert DER bytes as RawDer VALUE.
        let certs_content: Vec<u8> = signer_bufs
            .iter()
            .map(|b| b.cert_der.as_slice())
            .chain(self.extra_certs.iter().map(Vec::as_slice))
            .flat_map(|s| s.iter().copied())
            .collect();
        let certificates: Option<RawDer<'_>> = if certs_content.is_empty() {
            None
        } else {
            Some(RawDer(&certs_content))
        };

        let signed_data = SignedData {
            version: Integer::from(1),
            digest_algorithms: digest_algs,
            encap_content_info: eci,
            certificates,
            crls: None,
            signer_infos,
        };

        let sd_der = signed_data
            .to_der()
            .map_err(|e| PyValueError::new_err(format!("SignedData encode: {e}")))?;

        // ── Phase 4: ContentInfo wrapping ─────────────────────────────────────
        // ContentInfo ::= SEQUENCE { contentType OID, content [0] EXPLICIT ANY }
        // The ContentInfo generated type stores `content: RawDer` verbatim, so we
        // must include the [0] EXPLICIT tag+length in the bytes we hand to it.
        let sd_rawder = RawDer(&sd_der);
        let explicit0 = ExplicitTag::context_specific(0, &sd_rawder);
        let mut exp_enc = Encoder::new(Encoding::Der);
        explicit0
            .encode(&mut exp_enc)
            .map_err(|e| PyValueError::new_err(format!("explicit tag encode: {e}")))?;
        let content_bytes = exp_enc
            .finish()
            .map_err(|e| PyValueError::new_err(format!("{e}")))?;

        let id_signed_data =
            ObjectIdentifier::from_str("1.2.840.113549.1.7.2").expect("id-signedData OID is valid");
        let content_info = ContentInfo {
            content_type: id_signed_data,
            content: RawDer(&content_bytes),
        };

        let ci_der = content_info
            .to_der()
            .map_err(|e| PyValueError::new_err(format!("ContentInfo encode: {e}")))?;

        Ok(PyBytes::new(py, &ci_der))
    }
}