synta-python 0.3.0

//! Python bindings for CMS-KEM types (RFC 9629):
//! [`PyKEMRecipientInfo`] and [`PyCMSORIforKEMOtherInfo`].

use std::sync::OnceLock;

use pyo3::prelude::*;
use pyo3::types::PyBytes;

use synta::{Decoder, Encoding};

use crate::types::PyObjectIdentifier;

// ── PyKEMRecipientInfo ────────────────────────────────────────────────────────

/// KEM Recipient Info (RFC 9629) accessible from Python.
///
/// ``KEMRecipientInfo`` carries a quantum-safe KEM-encapsulated content-encryption
/// key inside a CMS ``EnvelopedData`` structure.  It is carried as an
/// ``OtherRecipientInfo`` alternative identified by ``id-ori-kem``.
///
/// ```python,ignore
/// kemri = KEMRecipientInfo.from_der(raw)
/// print(kemri.kem_algorithm_oid)  # e.g. ObjectIdentifier("1.3.6.1.4.1.22554.5.6.1")
/// key_bytes = kemri.encrypted_key
/// ```
#[pyclass(frozen, name = "KEMRecipientInfo")]
pub struct PyKEMRecipientInfo {
    _data: Py<PyBytes>,
    raw: &'static [u8],
    inner: OnceLock<Box<synta_certificate::cms_kem_types::KEMRecipientInfo<'static>>>,
    // OID caches
    kem_algorithm_oid_cache: OnceLock<Py<PyObjectIdentifier>>,
    kdf_algorithm_oid_cache: OnceLock<Py<PyObjectIdentifier>>,
    key_encryption_algorithm_oid_cache: OnceLock<Py<PyObjectIdentifier>>,
    // Optional DER params caches
    kem_algorithm_params_cache: OnceLock<Option<Py<PyBytes>>>,
    kdf_algorithm_params_cache: OnceLock<Option<Py<PyBytes>>>,
    key_encryption_algorithm_params_cache: OnceLock<Option<Py<PyBytes>>>,
    // Byte field caches
    recipient_id_cache: OnceLock<Py<PyBytes>>,
    kem_ciphertext_cache: OnceLock<Py<PyBytes>>,
    encrypted_key_cache: OnceLock<Py<PyBytes>>,
    // Optional byte field caches
    ukm_cache: OnceLock<Option<Py<PyBytes>>>,
}

impl PyKEMRecipientInfo {
    fn kemri(&self) -> PyResult<&synta_certificate::cms_kem_types::KEMRecipientInfo<'static>> {
        if let Some(v) = self.inner.get() {
            return Ok(v.as_ref());
        }
        let mut decoder = Decoder::new(self.raw, Encoding::Der);
        let decoded = decoder.decode().map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(format!(
                "KEMRecipientInfo DER decode failed: {e}"
            ))
        })?;
        let _ = self.inner.set(Box::new(decoded));
        Ok(self.inner.get().unwrap().as_ref())
    }
}

#[pymethods]
impl PyKEMRecipientInfo {
    /// Parse a DER-encoded ``KEMRecipientInfo`` SEQUENCE.
    #[staticmethod]
    fn from_der(py: Python<'_>, data: Bound<'_, PyBytes>) -> PyResult<Self> {
        let py_bytes = data.unbind();
        // SAFETY: `py_bytes` holds a strong reference (Py<PyBytes>) that
        // keeps the Python bytes object alive for the lifetime of this struct.
        // CPython's bytes objects have a fixed-address, non-relocating payload
        // buffer (CPython has no moving GC).  The slice lifetime is extended
        // to 'static; the actual safety invariants are:
        //   (1) All reads of `raw` go through `&self`; no borrow of the struct
        //       can outlive the struct, so `raw` is never read after drop begins.
        //   (2) `raw: &'static [u8]` has no destructor (fat pointer, no heap
        //       allocation), so field drop order does not cause use-after-free.
        //   (3) `inner` contains only borrow-typed fields in the decoded type;
        //       dropping the Box does not read through the contained &'static
        //       slices (borrows have no destructors in Rust).
        // CPython-only: does not hold for PyPy or GraalPy.
        let raw: &'static [u8] = unsafe {
            let s = py_bytes.bind(py).as_bytes();
            std::slice::from_raw_parts(s.as_ptr(), s.len())
        };
        {
            let mut d = Decoder::new(raw, Encoding::Der);
            d.read_tag()
                .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
            d.read_length()
                .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
        }
        Ok(Self {
            _data: py_bytes,
            raw,
            inner: OnceLock::new(),
            kem_algorithm_oid_cache: OnceLock::new(),
            kdf_algorithm_oid_cache: OnceLock::new(),
            key_encryption_algorithm_oid_cache: OnceLock::new(),
            kem_algorithm_params_cache: OnceLock::new(),
            kdf_algorithm_params_cache: OnceLock::new(),
            key_encryption_algorithm_params_cache: OnceLock::new(),
            recipient_id_cache: OnceLock::new(),
            kem_ciphertext_cache: OnceLock::new(),
            encrypted_key_cache: OnceLock::new(),
            ukm_cache: OnceLock::new(),
        })
    }

    /// Complete DER encoding of this ``KEMRecipientInfo``
    /// (the original bytes passed to ``from_der``).
    fn to_der<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        self._data.clone_ref(py).into_bound(py)
    }

    /// CMS version (always ``0`` for KEMRecipientInfo per RFC 9629).
    #[getter]
    fn version(&self) -> PyResult<i64> {
        Ok(self.kemri()?.version.as_i64().unwrap_or(0))
    }

    /// Raw DER bytes of the ``RecipientIdentifier`` CHOICE field.
    ///
    /// Decode with :class:`Decoder` to distinguish ``issuerAndSerialNumber``
    /// (SEQUENCE) from ``subjectKeyIdentifier`` (context tag ``[0]``).
    #[getter]
    fn recipient_id<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        if let Some(cached) = self.recipient_id_cache.get() {
            return Ok(cached.clone_ref(py).into_bound(py));
        }
        let py_bytes = PyBytes::new(py, self.kemri()?.rid.as_bytes()).unbind();
        let _ = self.recipient_id_cache.set(py_bytes.clone_ref(py));
        Ok(py_bytes.into_bound(py))
    }

    /// OID of the KEM algorithm (the ``kem`` field).
    ///
    /// For ML-KEM-768 this is ``2.16.840.1.101.3.4.4.2``.
    #[getter]
    fn kem_algorithm_oid<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyObjectIdentifier>> {
        if let Some(cached) = self.kem_algorithm_oid_cache.get() {
            return Ok(cached.clone_ref(py).into_bound(py));
        }
        let obj = Py::new(
            py,
            PyObjectIdentifier::from_oid(self.kemri()?.kem.algorithm.clone()),
        )?;
        let _ = self.kem_algorithm_oid_cache.set(obj.clone_ref(py));
        Ok(obj.into_bound(py))
    }

    /// Raw DER bytes of the KEM algorithm parameters, or ``None``.
    #[getter]
    fn kem_algorithm_params<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyBytes>>> {
        if let Some(cached) = self.kem_algorithm_params_cache.get() {
            return Ok(cached.as_ref().map(|b| b.clone_ref(py).into_bound(py)));
        }
        let computed = super::encode_element_opt(py, self.kemri()?.kem.parameters.as_ref())?;
        let to_store = computed.as_ref().map(|b| b.as_unbound().clone_ref(py));
        let _ = self.kem_algorithm_params_cache.set(to_store);
        Ok(computed)
    }

    /// Raw bytes of the KEM ciphertext (``kemct`` field, the encapsulated key).
    #[getter]
    fn kem_ciphertext<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        if let Some(cached) = self.kem_ciphertext_cache.get() {
            return Ok(cached.clone_ref(py).into_bound(py));
        }
        let py_bytes = PyBytes::new(py, self.kemri()?.kemct.as_bytes()).unbind();
        let _ = self.kem_ciphertext_cache.set(py_bytes.clone_ref(py));
        Ok(py_bytes.into_bound(py))
    }

    /// OID of the key-derivation algorithm (the ``kdf`` field).
    #[getter]
    fn kdf_algorithm_oid<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyObjectIdentifier>> {
        if let Some(cached) = self.kdf_algorithm_oid_cache.get() {
            return Ok(cached.clone_ref(py).into_bound(py));
        }
        let obj = Py::new(
            py,
            PyObjectIdentifier::from_oid(self.kemri()?.kdf.algorithm.clone()),
        )?;
        let _ = self.kdf_algorithm_oid_cache.set(obj.clone_ref(py));
        Ok(obj.into_bound(py))
    }

    /// Raw DER bytes of the KDF algorithm parameters, or ``None``.
    #[getter]
    fn kdf_algorithm_params<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyBytes>>> {
        if let Some(cached) = self.kdf_algorithm_params_cache.get() {
            return Ok(cached.as_ref().map(|b| b.clone_ref(py).into_bound(py)));
        }
        let computed = super::encode_element_opt(py, self.kemri()?.kdf.parameters.as_ref())?;
        let to_store = computed.as_ref().map(|b| b.as_unbound().clone_ref(py));
        let _ = self.kdf_algorithm_params_cache.set(to_store);
        Ok(computed)
    }

    /// KEK (key-encryption key) length in bytes (``kekLength`` field, range 1..65535).
    #[getter]
    fn kek_length(&self) -> PyResult<i64> {
        Ok(self.kemri()?.kek_length.as_i64().unwrap_or(0))
    }

    /// User keying material bytes (``ukm`` field), or ``None`` if absent.
    #[getter]
    fn ukm<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyBytes>>> {
        if let Some(cached) = self.ukm_cache.get() {
            return Ok(cached.as_ref().map(|b| b.clone_ref(py).into_bound(py)));
        }
        let computed = self
            .kemri()?
            .ukm
            .as_ref()
            .map(|u| PyBytes::new(py, u.as_bytes()));
        let to_store = computed.as_ref().map(|b| b.as_unbound().clone_ref(py));
        let _ = self.ukm_cache.set(to_store);
        Ok(computed)
    }

    /// OID of the key-encryption (key-wrap) algorithm (the ``wrap`` field).
    #[getter]
    fn key_encryption_algorithm_oid<'py>(
        &self,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, PyObjectIdentifier>> {
        if let Some(cached) = self.key_encryption_algorithm_oid_cache.get() {
            return Ok(cached.clone_ref(py).into_bound(py));
        }
        let obj = Py::new(
            py,
            PyObjectIdentifier::from_oid(self.kemri()?.wrap.algorithm.clone()),
        )?;
        let _ = self
            .key_encryption_algorithm_oid_cache
            .set(obj.clone_ref(py));
        Ok(obj.into_bound(py))
    }

    /// Raw DER bytes of the key-encryption algorithm parameters, or ``None``.
    #[getter]
    fn key_encryption_algorithm_params<'py>(
        &self,
        py: Python<'py>,
    ) -> PyResult<Option<Bound<'py, PyBytes>>> {
        if let Some(cached) = self.key_encryption_algorithm_params_cache.get() {
            return Ok(cached.as_ref().map(|b| b.clone_ref(py).into_bound(py)));
        }
        let computed = super::encode_element_opt(py, self.kemri()?.wrap.parameters.as_ref())?;
        let to_store = computed.as_ref().map(|b| b.as_unbound().clone_ref(py));
        let _ = self.key_encryption_algorithm_params_cache.set(to_store);
        Ok(computed)
    }

    /// Encrypted content-encryption key bytes (the ``encryptedKey`` OCTET STRING).
    #[getter]
    fn encrypted_key<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        if let Some(cached) = self.encrypted_key_cache.get() {
            return Ok(cached.clone_ref(py).into_bound(py));
        }
        let py_bytes = PyBytes::new(py, self.kemri()?.encrypted_key.as_bytes()).unbind();
        let _ = self.encrypted_key_cache.set(py_bytes.clone_ref(py));
        Ok(py_bytes.into_bound(py))
    }

    fn __repr__(&self) -> PyResult<String> {
        let kemri = self.kemri()?;
        Ok(format!(
            "KEMRecipientInfo(kem={}, kdf={}, kek_length={})",
            kemri.kem.algorithm,
            kemri.kdf.algorithm,
            kemri.kek_length.as_i64().unwrap_or(0),
        ))
    }
}

// ── PyCMSORIforKEMOtherInfo ───────────────────────────────────────────────────

/// CMS ORI KEM Other Info (RFC 9629 §6.2) accessible from Python.
///
/// ``CMSORIforKEMOtherInfo`` is used as the ``otherInfo`` input to the KDF
/// when deriving a KEK from a KEM shared secret.  It binds the key-encryption
/// algorithm, KEK length, and optional UKM to the KDF computation.
///
/// ```python,ignore
/// info = CMSORIforKEMOtherInfo.from_der(raw)
/// print(info.kek_length)      # e.g. 32
/// print(info.key_encryption_algorithm_oid)
/// ```
#[pyclass(frozen, name = "CMSORIforKEMOtherInfo")]
pub struct PyCMSORIforKEMOtherInfo {
    _data: Py<PyBytes>,
    raw: &'static [u8],
    inner: OnceLock<Box<synta_certificate::cms_kem_types::CMSORIforKEMOtherInfo<'static>>>,
    key_encryption_algorithm_oid_cache: OnceLock<Py<PyObjectIdentifier>>,
    key_encryption_algorithm_params_cache: OnceLock<Option<Py<PyBytes>>>,
    ukm_cache: OnceLock<Option<Py<PyBytes>>>,
}

impl PyCMSORIforKEMOtherInfo {
    fn info(&self) -> PyResult<&synta_certificate::cms_kem_types::CMSORIforKEMOtherInfo<'static>> {
        if let Some(v) = self.inner.get() {
            return Ok(v.as_ref());
        }
        let mut decoder = Decoder::new(self.raw, Encoding::Der);
        let decoded = decoder.decode().map_err(|e| {
            pyo3::exceptions::PyValueError::new_err(format!(
                "CMSORIforKEMOtherInfo DER decode failed: {e}"
            ))
        })?;
        let _ = self.inner.set(Box::new(decoded));
        Ok(self.inner.get().unwrap().as_ref())
    }
}

#[pymethods]
impl PyCMSORIforKEMOtherInfo {
    /// Parse a DER-encoded ``CMSORIforKEMOtherInfo`` SEQUENCE.
    #[staticmethod]
    fn from_der(py: Python<'_>, data: Bound<'_, PyBytes>) -> PyResult<Self> {
        let py_bytes = data.unbind();
        // SAFETY: `py_bytes` holds a strong reference (Py<PyBytes>) that
        // keeps the Python bytes object alive for the lifetime of this struct.
        // CPython's bytes objects have a fixed-address, non-relocating payload
        // buffer (CPython has no moving GC).  The slice lifetime is extended
        // to 'static; the actual safety invariants are:
        //   (1) All reads of `raw` go through `&self`; no borrow of the struct
        //       can outlive the struct, so `raw` is never read after drop begins.
        //   (2) `raw: &'static [u8]` has no destructor (fat pointer, no heap
        //       allocation), so field drop order does not cause use-after-free.
        //   (3) `inner` contains only borrow-typed fields in the decoded type;
        //       dropping the Box does not read through the contained &'static
        //       slices (borrows have no destructors in Rust).
        // CPython-only: does not hold for PyPy or GraalPy.
        let raw: &'static [u8] = unsafe {
            let s = py_bytes.bind(py).as_bytes();
            std::slice::from_raw_parts(s.as_ptr(), s.len())
        };
        {
            let mut d = Decoder::new(raw, Encoding::Der);
            d.read_tag()
                .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
            d.read_length()
                .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
        }
        Ok(Self {
            _data: py_bytes,
            raw,
            inner: OnceLock::new(),
            key_encryption_algorithm_oid_cache: OnceLock::new(),
            key_encryption_algorithm_params_cache: OnceLock::new(),
            ukm_cache: OnceLock::new(),
        })
    }

    /// Complete DER encoding of this ``CMSORIforKEMOtherInfo``
    /// (the original bytes passed to ``from_der``).
    fn to_der<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        self._data.clone_ref(py).into_bound(py)
    }

    /// OID of the key-encryption (key-wrap) algorithm.
    #[getter]
    fn key_encryption_algorithm_oid<'py>(
        &self,
        py: Python<'py>,
    ) -> PyResult<Bound<'py, PyObjectIdentifier>> {
        if let Some(cached) = self.key_encryption_algorithm_oid_cache.get() {
            return Ok(cached.clone_ref(py).into_bound(py));
        }
        let obj = Py::new(
            py,
            PyObjectIdentifier::from_oid(self.info()?.wrap.algorithm.clone()),
        )?;
        let _ = self
            .key_encryption_algorithm_oid_cache
            .set(obj.clone_ref(py));
        Ok(obj.into_bound(py))
    }

    /// Raw DER bytes of the key-encryption algorithm parameters, or ``None``.
    #[getter]
    fn key_encryption_algorithm_params<'py>(
        &self,
        py: Python<'py>,
    ) -> PyResult<Option<Bound<'py, PyBytes>>> {
        if let Some(cached) = self.key_encryption_algorithm_params_cache.get() {
            return Ok(cached.as_ref().map(|b| b.clone_ref(py).into_bound(py)));
        }
        let computed = super::encode_element_opt(py, self.info()?.wrap.parameters.as_ref())?;
        let to_store = computed.as_ref().map(|b| b.as_unbound().clone_ref(py));
        let _ = self.key_encryption_algorithm_params_cache.set(to_store);
        Ok(computed)
    }

    /// KEK length in bytes (``kekLength`` field, range 1..65535).
    #[getter]
    fn kek_length(&self) -> PyResult<i64> {
        Ok(self.info()?.kek_length.as_i64().unwrap_or(0))
    }

    /// User keying material bytes (``ukm`` field), or ``None`` if absent.
    #[getter]
    fn ukm<'py>(&self, py: Python<'py>) -> PyResult<Option<Bound<'py, PyBytes>>> {
        if let Some(cached) = self.ukm_cache.get() {
            return Ok(cached.as_ref().map(|b| b.clone_ref(py).into_bound(py)));
        }
        let computed = self
            .info()?
            .ukm
            .as_ref()
            .map(|u| PyBytes::new(py, u.as_bytes()));
        let to_store = computed.as_ref().map(|b| b.as_unbound().clone_ref(py));
        let _ = self.ukm_cache.set(to_store);
        Ok(computed)
    }

    fn __repr__(&self) -> PyResult<String> {
        let info = self.info()?;
        Ok(format!(
            "CMSORIforKEMOtherInfo(wrap={}, kek_length={})",
            info.wrap.algorithm,
            info.kek_length.as_i64().unwrap_or(0),
        ))
    }
}