synta-python 0.1.3

//! Python bindings for RFC 9810 Certificate Management Protocol (CMP) types.
//!
//! Exposes ``CMPMessage`` (wrapping ``PKIMessage``) as a Python class and
//! installs MAC algorithm and key-purpose OID constants into the ``synta.cmp``
//! submodule.

use std::sync::OnceLock;

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

use synta::traits::Encode;
use synta::{Decoder, Encoding};

use crate::error::SyntaErr;
use crate::types::PyObjectIdentifier;

// ── helpers ───────────────────────────────────────────────────────────────────

fn encode_to_der<T: Encode>(v: &T) -> Vec<u8> {
    let mut enc = synta::Encoder::new(Encoding::Der);
    if v.encode(&mut enc).is_err() {
        return Vec::new();
    }
    enc.finish().unwrap_or_default()
}

/// Map a `PKIBody` variant to its CMP message-type name (lowercase, per RFC 9810).
fn body_type_name(body: &synta_certificate::cmp_types::PKIBody<'_>) -> &'static str {
    use synta_certificate::cmp_types::PKIBody::*;
    match body {
        Ir(_) => "ir",
        Ip(_) => "ip",
        Cr(_) => "cr",
        Cp(_) => "cp",
        P10cr(_) => "p10cr",
        Popdecc(_) => "popdecc",
        Popdecr(_) => "popdecr",
        Kur(_) => "kur",
        Kup(_) => "kup",
        Krr(_) => "krr",
        Krp(_) => "krp",
        Rr(_) => "rr",
        Rp(_) => "rp",
        Ccr(_) => "ccr",
        Ccp(_) => "ccp",
        Ckuann(_) => "ckuann",
        Cann(_) => "cann",
        Rann(_) => "rann",
        Crlann(_) => "crlann",
        Pkiconf(_) => "pkiconf",
        Nested(_) => "nested",
        Genm(_) => "genm",
        Genp(_) => "genp",
        Error(_) => "error",
        CertConf(_) => "certConf",
        PollReq(_) => "pollReq",
        PollRep(_) => "pollRep",
    }
}

/// Extract raw body DER from a `PKIBody` variant.  Returns `None` for
/// the `pkiconf` (NULL) arm; returns the `RawDer` bytes for all others.
fn body_raw_der(body: &synta_certificate::cmp_types::PKIBody<'_>) -> Option<Vec<u8>> {
    use synta_certificate::cmp_types::PKIBody::*;
    match body {
        Pkiconf(_) => None,
        Ir(r) | Ip(r) | Cr(r) | Cp(r) | P10cr(r) | Popdecc(r) | Popdecr(r) | Kur(r) | Kup(r)
        | Krr(r) | Krp(r) | Rr(r) | Rp(r) | Ccr(r) | Ccp(r) | Ckuann(r) | Cann(r) | Rann(r)
        | Crlann(r) | Nested(r) | Genm(r) | Genp(r) | Error(r) | CertConf(r) | PollReq(r)
        | PollRep(r) => Some(r.as_bytes().to_vec()),
    }
}

// ── PyCMPMessage ──────────────────────────────────────────────────────────────

/// A CMP ``PKIMessage`` (RFC 9810 §5.1).
///
/// ``PKIMessage`` is the top-level CMP envelope, carrying a ``PKIHeader``
/// (metadata: pvno, sender, recipient, timestamps, nonces, …) and a
/// ``PKIBody`` (the message payload: IR, IP, CR, CP, RR, RP, …).
///
/// The ``body_type`` property names the active ``PKIBody`` arm in lowercase
/// (``"ir"``, ``"ip"``, ``"cr"``, ``"cp"``, ``"pkiconf"``, …).  For all arms
/// except ``pkiconf``, ``body_der`` returns the raw DER bytes of the body
/// content for further decoding with a ``synta.Decoder``.
///
/// ```python,ignore
/// import synta.cmp as cmp
///
/// msg = cmp.CMPMessage.from_der(open("cmp.der", "rb").read())
/// print(msg.pvno, msg.body_type)
/// if msg.body_type == "ir":
///     # body_der contains the DER of CertReqMessages
///     import synta.crmf as crmf
///     reqs = crmf.CertReqMessages.from_der(msg.body_der)
///     for req in reqs:
///         print(req.cert_req_id, req.popo_type)
/// ```
#[pyclass(frozen, name = "CMPMessage")]
pub struct PyCMPMessage {
    _data: Py<PyBytes>,
    raw: &'static [u8],
    inner: OnceLock<Box<synta_certificate::cmp_types::PKIMessage<'static>>>,
    // Field caches
    pvno_cache: OnceLock<i64>,
    body_type_cache: OnceLock<&'static str>,
    body_der_cache: OnceLock<Option<Py<PyBytes>>>,
    sender_der_cache: OnceLock<Py<PyBytes>>,
    recipient_der_cache: OnceLock<Py<PyBytes>>,
    transaction_id_cache: OnceLock<Option<Py<PyBytes>>>,
    sender_nonce_cache: OnceLock<Option<Py<PyBytes>>>,
    recip_nonce_cache: OnceLock<Option<Py<PyBytes>>>,
    protection_alg_oid_cache: OnceLock<Option<Py<PyObjectIdentifier>>>,
    message_time_cache: OnceLock<Option<Py<PyString>>>,
}

impl PyCMPMessage {
    fn msg(&self) -> PyResult<&synta_certificate::cmp_types::PKIMessage<'static>> {
        if let Some(v) = self.inner.get() {
            return Ok(v.as_ref());
        }
        let mut dec = Decoder::new(self.raw, Encoding::Der);
        let decoded = dec
            .decode::<synta_certificate::cmp_types::PKIMessage<'_>>()
            .map_err(SyntaErr)?;
        // SAFETY: raw is pinned by _data for the lifetime of self.
        let decoded: synta_certificate::cmp_types::PKIMessage<'static> =
            unsafe { std::mem::transmute(decoded) };
        let _ = self.inner.set(Box::new(decoded));
        Ok(self.inner.get().unwrap().as_ref())
    }
}

#[pymethods]
impl PyCMPMessage {
    /// Parse a DER-encoded ``PKIMessage`` SEQUENCE.
    ///
    /// :param data: DER bytes of the ``PKIMessage``.
    /// :raises ValueError: if the bytes cannot be decoded.
    #[staticmethod]
    fn from_der(py: Python<'_>, data: Bound<'_, PyBytes>) -> PyResult<Self> {
        let py_bytes = data.unbind();
        {
            let raw = py_bytes.as_bytes(py);
            Decoder::new(raw, Encoding::Der)
                .decode::<synta_certificate::cmp_types::PKIMessage<'_>>()
                .map_err(SyntaErr)?;
        }
        let raw: &'static [u8] = unsafe { std::mem::transmute(py_bytes.as_bytes(py)) };
        Ok(Self {
            _data: py_bytes,
            raw,
            inner: OnceLock::new(),
            pvno_cache: OnceLock::new(),
            body_type_cache: OnceLock::new(),
            body_der_cache: OnceLock::new(),
            sender_der_cache: OnceLock::new(),
            recipient_der_cache: OnceLock::new(),
            transaction_id_cache: OnceLock::new(),
            sender_nonce_cache: OnceLock::new(),
            recip_nonce_cache: OnceLock::new(),
            protection_alg_oid_cache: OnceLock::new(),
            message_time_cache: OnceLock::new(),
        })
    }

    /// Return the DER encoding of this ``PKIMessage``.
    fn to_der<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        let mut enc = synta::Encoder::new(Encoding::Der);
        self.msg()?.encode(&mut enc).map_err(SyntaErr)?;
        Ok(PyBytes::new(py, &enc.finish().map_err(SyntaErr)?))
    }

    /// CMP protocol version number (``pvno`` field; ``2`` for CMP v2/v3).
    #[getter]
    fn pvno(&self) -> PyResult<i64> {
        if let Some(v) = self.pvno_cache.get() {
            return Ok(*v);
        }
        let v = self.msg()?.header.pvno.as_i64().unwrap_or(2);
        let _ = self.pvno_cache.set(v);
        Ok(v)
    }

    /// Name of the active ``PKIBody`` arm in lowercase (e.g. ``"ir"``, ``"ip"``,
    /// ``"cr"``, ``"cp"``, ``"rr"``, ``"rp"``, ``"pkiconf"``).
    #[getter]
    fn body_type(&self) -> PyResult<&'static str> {
        if let Some(v) = self.body_type_cache.get() {
            return Ok(v);
        }
        let name = body_type_name(&self.msg()?.body);
        let _ = self.body_type_cache.set(name);
        Ok(name)
    }

    /// Raw DER bytes of the ``PKIBody`` content, or ``None`` for ``pkiconf``
    /// (which carries a NULL value and has no further content to decode).
    ///
    /// For ``"ir"``, ``"cr"`` etc., these bytes are a ``CertReqMessages``
    /// SEQUENCE OF that can be decoded with ``synta.crmf.CertReqMessages.from_der``.
    #[getter]
    fn body_der<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyBytes>>> {
        if let Some(cached) = self.body_der_cache.get() {
            return Ok(cached.as_ref().map(|b| b.clone_ref(py).into_bound(py)));
        }
        let der = body_raw_der(&self.msg()?.body);
        let py_opt = der
            .as_deref()
            .map(|b| PyBytes::new(py, b).as_unbound().clone_ref(py));
        let _ = self.body_der_cache.set(py_opt);
        Ok(self
            .body_der_cache
            .get()
            .unwrap()
            .as_ref()
            .map(|b| b.clone_ref(py).into_bound(py)))
    }

    /// Raw DER bytes of the ``sender`` ``GeneralName`` field.
    #[getter]
    fn sender_der<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        if let Some(c) = self.sender_der_cache.get() {
            return Ok(c.clone_ref(py).into_bound(py));
        }
        let der = encode_to_der(&self.msg()?.header.sender);
        let b = PyBytes::new(py, &der);
        let _ = self.sender_der_cache.set(b.as_unbound().clone_ref(py));
        Ok(b)
    }

    /// Raw DER bytes of the ``recipient`` ``GeneralName`` field.
    #[getter]
    fn recipient_der<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        if let Some(c) = self.recipient_der_cache.get() {
            return Ok(c.clone_ref(py).into_bound(py));
        }
        let der = encode_to_der(&self.msg()?.header.recipient);
        let b = PyBytes::new(py, &der);
        let _ = self.recipient_der_cache.set(b.as_unbound().clone_ref(py));
        Ok(b)
    }

    /// ``transactionID`` bytes, or ``None`` if not present.
    #[getter]
    fn transaction_id<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyBytes>>> {
        if let Some(cached) = self.transaction_id_cache.get() {
            return Ok(cached.as_ref().map(|b| b.clone_ref(py).into_bound(py)));
        }
        let opt = self
            .msg()?
            .header
            .transaction_id
            .as_ref()
            .map(|o| PyBytes::new(py, o.as_bytes()).as_unbound().clone_ref(py));
        let _ = self.transaction_id_cache.set(opt);
        Ok(self
            .transaction_id_cache
            .get()
            .unwrap()
            .as_ref()
            .map(|b| b.clone_ref(py).into_bound(py)))
    }

    /// ``senderNonce`` bytes, or ``None`` if not present.
    #[getter]
    fn sender_nonce<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyBytes>>> {
        if let Some(cached) = self.sender_nonce_cache.get() {
            return Ok(cached.as_ref().map(|b| b.clone_ref(py).into_bound(py)));
        }
        let opt = self
            .msg()?
            .header
            .sender_nonce
            .as_ref()
            .map(|o| PyBytes::new(py, o.as_bytes()).as_unbound().clone_ref(py));
        let _ = self.sender_nonce_cache.set(opt);
        Ok(self
            .sender_nonce_cache
            .get()
            .unwrap()
            .as_ref()
            .map(|b| b.clone_ref(py).into_bound(py)))
    }

    /// ``recipNonce`` bytes, or ``None`` if not present.
    #[getter]
    fn recip_nonce<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyBytes>>> {
        if let Some(cached) = self.recip_nonce_cache.get() {
            return Ok(cached.as_ref().map(|b| b.clone_ref(py).into_bound(py)));
        }
        let opt = self
            .msg()?
            .header
            .recip_nonce
            .as_ref()
            .map(|o| PyBytes::new(py, o.as_bytes()).as_unbound().clone_ref(py));
        let _ = self.recip_nonce_cache.set(opt);
        Ok(self
            .recip_nonce_cache
            .get()
            .unwrap()
            .as_ref()
            .map(|b| b.clone_ref(py).into_bound(py)))
    }

    /// OID of the ``protectionAlg`` ``AlgorithmIdentifier``, or ``None``.
    #[getter]
    fn protection_alg_oid(&self, py: Python<'_>) -> PyResult<Option<Py<PyObjectIdentifier>>> {
        if let Some(cached) = self.protection_alg_oid_cache.get() {
            return Ok(cached.as_ref().map(|o| o.clone_ref(py)));
        }
        let opt = self
            .msg()?
            .header
            .protection_alg
            .as_ref()
            .map(|alg| Py::new(py, PyObjectIdentifier::from_oid(alg.algorithm.clone())))
            .transpose()?;
        let _ = self
            .protection_alg_oid_cache
            .set(opt.as_ref().map(|o| o.clone_ref(py)));
        Ok(opt)
    }

    /// ``messageTime`` as a GeneralizedTime string, or ``None`` if not present.
    #[getter]
    fn message_time<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyString>>> {
        if let Some(cached) = self.message_time_cache.get() {
            return Ok(cached.as_ref().map(|s| s.clone_ref(py).into_bound(py)));
        }
        let opt = self
            .msg()?
            .header
            .message_time
            .as_ref()
            .map(|t| PyString::new(py, &t.to_string()).as_unbound().clone_ref(py));
        let _ = self.message_time_cache.set(opt);
        Ok(self
            .message_time_cache
            .get()
            .unwrap()
            .as_ref()
            .map(|s| s.clone_ref(py).into_bound(py)))
    }

    fn __repr__(&self) -> PyResult<String> {
        Ok(format!(
            "CMPMessage(pvno={}, body_type={})",
            self.pvno()?,
            self.body_type()?,
        ))
    }
}

// ── register_cmp_submodule ────────────────────────────────────────────────────

/// Build and register the ``synta.cmp`` submodule.
pub(super) fn register_cmp_submodule(parent: &Bound<'_, PyModule>) -> PyResult<()> {
    let py = parent.py();
    let m = PyModule::new(py, "cmp")?;

    m.add_class::<PyCMPMessage>()?;

    // ── MAC algorithm OIDs (RFC 9810 Appendix D) ─────────────────────────────
    m.add(
        "ID_PASSWORD_BASED_MAC",
        super::oid_const(py, synta_certificate::cmp_types::ID_PASSWORD_BASED_MAC),
    )?;
    m.add(
        "ID_DHBASED_MAC",
        super::oid_const(py, synta_certificate::cmp_types::ID_DHBASED_MAC),
    )?;
    m.add(
        "ID_KEM_BASED_MAC",
        super::oid_const(py, synta_certificate::cmp_types::ID_KEM_BASED_MAC),
    )?;

    // ── CMP key-purpose OID (RFC 9810 §4) ────────────────────────────────────
    m.add(
        "ID_KP_CM_KGA",
        super::oid_const(py, synta_certificate::cmp_types::ID_KP_CM_KGA),
    )?;

    // ── CMP registration-control OIDs shared with CRMF (RFC 9810 Appendix) ──
    m.add(
        "ID_REG_CTRL_ALT_CERT_TEMPLATE",
        super::oid_const(
            py,
            synta_certificate::cmp_types::ID_REG_CTRL_ALT_CERT_TEMPLATE,
        ),
    )?;
    m.add(
        "ID_REG_CTRL_ALG_ID",
        super::oid_const(py, synta_certificate::cmp_types::ID_REG_CTRL_ALG_ID),
    )?;
    m.add(
        "ID_REG_CTRL_RSA_KEY_LEN",
        super::oid_const(py, synta_certificate::cmp_types::ID_REG_CTRL_RSA_KEY_LEN),
    )?;

    crate::install_submodule(
        parent,
        &m,
        "synta.cmp",
        Some(concat!(
            "synta.cmp — RFC 9810 Certificate Management Protocol v3 types.\n\n",
            "Provides CMPMessage (wrapping PKIMessage) for decoding CMP\n",
            "envelopes.  body_type names the active PKIBody arm; body_der\n",
            "returns the raw bytes for further decoding with synta.crmf or\n",
            "a synta.Decoder.  Also exposes MAC algorithm and key-purpose\n",
            "OID constants.",
        )),
    )
}