synta-python-mtc 0.2.6

use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::PyBytes;
use synta::traits::{Decode, Encode};

use super::proof::PyHashAlgorithm;

// ── Module-level utility functions ────────────────────────────────────────────

/// Extract the leaf index from a composite MTC serial number (draft-04 §6.1).
///
/// The serial number encodes ``(log_number << 48) | log_entry_index``.
/// This function returns the lower 48 bits, which is the log entry index.
///
/// ```python,ignore
/// import synta.mtc as mtc
///
/// serial = (3 << 48) | 42  # log_number=3, leaf_index=42
/// assert mtc.extract_leaf_index(serial) == 42
/// ```
#[pyfunction]
pub fn extract_leaf_index(serial: u64) -> u64 {
    serial & 0xFFFF_FFFF_FFFF
}

/// Extract the log number from a composite MTC serial number (draft-04 §6.1).
///
/// The serial number encodes ``(log_number << 48) | log_entry_index``.
/// This function returns the upper 16 bits, which is the log number.
///
/// ```python,ignore
/// import synta.mtc as mtc
///
/// serial = (3 << 48) | 42  # log_number=3, leaf_index=42
/// assert mtc.extract_log_number(serial) == 3
/// ```
#[pyfunction]
pub fn extract_log_number(serial: u64) -> u64 {
    serial >> 48
}

/// Build the DER-encoded value for the ``id-pe-mtcCertificationAuthority``
/// extension (draft-04 §5.5).
///
/// The extension MUST be marked critical in certificates issued by an
/// MTC-capable CA.  It advertises the hash algorithm, signature algorithm,
/// and minimum serial number used by the CA's MTC logs.
///
/// :param hash_alg: :class:`HashAlgorithm` used by all CA logs.
/// :param sig_alg_der: DER-encoded ``AlgorithmIdentifier`` SEQUENCE TLV for
///     the CA cosigner signature algorithm.
/// :param min_serial: Minimum allowed serial number.
/// :returns: ``(extension_value_der, critical)`` where ``critical`` is
///     always ``True`` per spec §5.5.
/// :raises ValueError: if *sig_alg_der* cannot be decoded or encoding fails.
#[pyfunction]
pub fn build_mtc_ca_extension<'py>(
    py: Python<'py>,
    hash_alg: &PyHashAlgorithm,
    sig_alg_der: &[u8],
    min_serial: u64,
) -> PyResult<(Bound<'py, PyBytes>, bool)> {
    use synta_certificate::AlgorithmIdentifier;

    let mut dec = synta::Decoder::new(sig_alg_der, synta::Encoding::Der);
    let sig_alg = AlgorithmIdentifier::decode(&mut dec)
        .map_err(|e| PyValueError::new_err(format!("invalid AlgorithmIdentifier DER: {e}")))?;

    let (der, critical) = synta_mtc::builder::ca_extension::build_mtc_ca_extension_from_hash(
        hash_alg.inner,
        &sig_alg,
        min_serial,
    )
    .map_err(|e| PyValueError::new_err(e.to_string()))?;

    Ok((PyBytes::new(py, &der), critical))
}

/// Parse the ``id-pe-mtcCertificationAuthority`` extension from a
/// DER-encoded X.509 certificate (draft-04 §5.5).
///
/// Searches the certificate's extensions for
/// ``id-pe-mtcCertificationAuthority``.
///
/// :param cert_der: DER-encoded ``Certificate`` bytes (the full certificate,
///     not just the extension value).
/// :returns: ``(hash_algorithm, sig_alg_der, min_serial)`` where
///     ``hash_algorithm`` is a :class:`HashAlgorithm`, ``sig_alg_der`` is the
///     DER-encoded ``AlgorithmIdentifier`` for the CA cosigner signature
///     algorithm, and ``min_serial`` is the minimum allowed serial number.
///     Returns ``None`` if the extension is absent.
/// :raises ValueError: if *cert_der* is malformed or the extension value
///     cannot be decoded.
#[pyfunction]
#[allow(clippy::type_complexity)]
pub fn parse_mtc_ca_extension<'py>(
    py: Python<'py>,
    cert_der: &[u8],
) -> PyResult<Option<(Py<PyHashAlgorithm>, Bound<'py, PyBytes>, u64)>> {
    use synta::ObjectIdentifier;
    use synta_mtc::builder::ca_extension::parse_mtc_ca_extension as rust_parse;
    use synta_mtc::crypto::hash::HashAlgorithm;

    let parsed = rust_parse(cert_der).map_err(|e| PyValueError::new_err(e.to_string()))?;

    match parsed {
        None => Ok(None),
        Some(ext) => {
            // Resolve the hash algorithm from its OID components.
            let hash_alg =
                HashAlgorithm::from_oid_components(&ext.log_hash_oid).ok_or_else(|| {
                    let oid = ObjectIdentifier::new(&ext.log_hash_oid)
                        .map(|o| o.to_string())
                        .unwrap_or_else(|_| format!("{:?}", ext.log_hash_oid));
                    PyValueError::new_err(format!("unknown hash algorithm OID: {oid}"))
                })?;
            let py_hash_alg = Py::new(py, PyHashAlgorithm { inner: hash_alg })?;

            // Re-encode the signature algorithm OID components back to DER via
            // an AlgorithmIdentifier with no parameters.
            let sig_oid = ObjectIdentifier::new(&ext.sig_alg_oid).map_err(|e| {
                PyValueError::new_err(format!("invalid sig_alg OID components: {e}"))
            })?;
            let sig_alg = synta_certificate::AlgorithmIdentifier {
                algorithm: sig_oid,
                parameters: None,
            };
            let mut enc = synta::Encoder::new(synta::Encoding::Der);
            sig_alg
                .encode(&mut enc)
                .map_err(|e| PyValueError::new_err(format!("encode sig_alg: {e}")))?;
            let sig_alg_der = enc
                .finish()
                .map_err(|e| PyValueError::new_err(format!("finish sig_alg encoding: {e}")))?;

            Ok(Some((
                py_hash_alg,
                PyBytes::new(py, &sig_alg_der),
                ext.min_serial,
            )))
        }
    }
}