synta-python 0.2.5

//! ML-DSA, RSA public-key, and RSASSA-PSS parameter types for the Python binding.

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

use synta::{Decoder, Encoding};

use crate::types::PyObjectIdentifier;

// ── ML-DSA private key ────────────────────────────────────────────────────────

/// A parsed ML-DSA private key (RFC 9881 / FIPS 204).
///
/// ML-DSA private keys are DER-encoded ``CHOICE`` structures with three
/// possible representations:
///
/// - ``"seed"`` — only the 32-octet seed; the expanded key is rederived on use.
/// - ``"expanded_key"`` — only the pre-expanded key material; faster signing.
/// - ``"both"`` — seed and expanded key stored together.
///
/// The :attr:`parameter_set` property exposes the algorithm
/// :class:`~synta.ObjectIdentifier` (one of ``synta.oids.ML_DSA_44``,
/// ``synta.oids.ML_DSA_65``, or ``synta.oids.ML_DSA_87``).
///
/// Use :meth:`from_der` to decode the ``privateKey`` value bytes from a
/// ``OneAsymmetricKey`` / PKCS\#8 structure.  Pass the algorithm OID from the
/// ``OneAsymmetricKey.algorithm`` field so the correct parameter set is used
/// automatically.
#[pyclass(frozen, name = "MlDsaPrivateKey")]
pub struct PyMlDsaPrivateKey {
    parameter_set: Py<PyObjectIdentifier>,
    kind: &'static str,
    seed: Option<Py<PyBytes>>,
    expanded_key: Option<Py<PyBytes>>,
}

#[pymethods]
impl PyMlDsaPrivateKey {
    /// Decode an ML-DSA private key CHOICE from DER bytes.
    ///
    /// ``algorithm_oid`` selects the ML-DSA parameter set and must be one of
    /// ``synta.oids.ML_DSA_44``, ``synta.oids.ML_DSA_65``, or
    /// ``synta.oids.ML_DSA_87`` (or equivalent dotted-decimal strings).
    /// Pass the OID directly from the ``OneAsymmetricKey.algorithm`` field so
    /// the parameter set is discovered automatically without hard-coding numbers.
    ///
    /// ``data`` must be the raw ``privateKey`` value bytes from a
    /// ``OneAsymmetricKey`` / PKCS\#8 structure.
    ///
    /// :param algorithm_oid: :class:`~synta.ObjectIdentifier` or dotted-decimal
    ///     ``str``; must be one of the ``synta.oids.ML_DSA_*`` constants.
    /// :param data: DER bytes of the ``ML-DSA-{44,65,87}-PrivateKey`` CHOICE.
    /// :returns: :class:`MlDsaPrivateKey`
    /// :raises ValueError: on unknown OID or malformed DER.
    ///
    /// ```python,ignore
    /// # oid comes from parsing the OneAsymmetricKey structure:
    /// key = synta.MlDsaPrivateKey.from_der(oid, private_key_bytes)
    /// if key.kind == "both":
    ///     seed = key.seed
    ///     expanded = key.expanded_key
    /// ```
    #[staticmethod]
    fn from_der(
        py: Python<'_>,
        algorithm_oid: &Bound<'_, PyAny>,
        data: &Bound<'_, PyBytes>,
    ) -> PyResult<Self> {
        use synta_certificate::oids;
        use synta_certificate::{
            MlDsa44PrivateKey, MlDsa44PrivateKeyBoth, MlDsa65PrivateKey, MlDsa65PrivateKeyBoth,
            MlDsa87PrivateKey, MlDsa87PrivateKeyBoth,
        };

        let oid = super::oid_from_pyany(algorithm_oid)?;
        let comps = oid.components();
        let raw = data.as_bytes();
        let mut decoder = Decoder::new(raw, Encoding::Der);

        macro_rules! to_py {
            ($slice:expr) => {
                PyBytes::new(py, $slice).unbind()
            };
        }

        let (kind, seed, expanded_key) = if comps == oids::ML_DSA_44 {
            let key: MlDsa44PrivateKey<'_> = decoder
                .decode()
                .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
            match key {
                MlDsa44PrivateKey::Seed(s) => ("seed", Some(to_py!(s.as_bytes())), None),
                MlDsa44PrivateKey::ExpandedKey(ek) => {
                    ("expanded_key", None, Some(to_py!(ek.as_bytes())))
                }
                MlDsa44PrivateKey::Both(MlDsa44PrivateKeyBoth {
                    seed: s,
                    expanded_key: ek,
                }) => (
                    "both",
                    Some(to_py!(s.as_bytes())),
                    Some(to_py!(ek.as_bytes())),
                ),
            }
        } else if comps == oids::ML_DSA_65 {
            let key: MlDsa65PrivateKey<'_> = decoder
                .decode()
                .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
            match key {
                MlDsa65PrivateKey::Seed(s) => ("seed", Some(to_py!(s.as_bytes())), None),
                MlDsa65PrivateKey::ExpandedKey(ek) => {
                    ("expanded_key", None, Some(to_py!(ek.as_bytes())))
                }
                MlDsa65PrivateKey::Both(MlDsa65PrivateKeyBoth {
                    seed: s,
                    expanded_key: ek,
                }) => (
                    "both",
                    Some(to_py!(s.as_bytes())),
                    Some(to_py!(ek.as_bytes())),
                ),
            }
        } else if comps == oids::ML_DSA_87 {
            let key: MlDsa87PrivateKey<'_> = decoder
                .decode()
                .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
            match key {
                MlDsa87PrivateKey::Seed(s) => ("seed", Some(to_py!(s.as_bytes())), None),
                MlDsa87PrivateKey::ExpandedKey(ek) => {
                    ("expanded_key", None, Some(to_py!(ek.as_bytes())))
                }
                MlDsa87PrivateKey::Both(MlDsa87PrivateKeyBoth {
                    seed: s,
                    expanded_key: ek,
                }) => (
                    "both",
                    Some(to_py!(s.as_bytes())),
                    Some(to_py!(ek.as_bytes())),
                ),
            }
        } else {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "unknown ML-DSA algorithm OID: {oid}; expected one of ML_DSA_44, ML_DSA_65, ML_DSA_87",
            )));
        };

        let ps = Py::new(py, PyObjectIdentifier::from_oid(oid))
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
        Ok(Self {
            parameter_set: ps,
            kind,
            seed,
            expanded_key,
        })
    }

    /// The algorithm :class:`~synta.ObjectIdentifier` identifying the ML-DSA
    /// parameter set (one of ``synta.oids.ML_DSA_44``, ``synta.oids.ML_DSA_65``,
    /// or ``synta.oids.ML_DSA_87``).
    #[getter]
    fn parameter_set<'py>(&self, py: Python<'py>) -> Bound<'py, PyObjectIdentifier> {
        self.parameter_set.bind(py).clone()
    }

    /// Which representation is stored: ``"seed"``, ``"expanded_key"``, or
    /// ``"both"``.
    #[getter]
    fn kind(&self) -> &str {
        self.kind
    }

    /// The 32-octet seed bytes, or ``None`` when :attr:`kind` is
    /// ``"expanded_key"``.
    #[getter]
    fn seed<'py>(&self, py: Python<'py>) -> Option<Bound<'py, PyBytes>> {
        self.seed.as_ref().map(|b| b.bind(py).clone())
    }

    /// The expanded key bytes, or ``None`` when :attr:`kind` is ``"seed"``.
    #[getter]
    fn expanded_key<'py>(&self, py: Python<'py>) -> Option<Bound<'py, PyBytes>> {
        self.expanded_key.as_ref().map(|b| b.bind(py).clone())
    }

    fn __repr__(&self, py: Python<'_>) -> String {
        let oid_str = self.parameter_set.bind(py).borrow().inner.to_string();
        format!(
            "MlDsaPrivateKey(parameter_set={oid_str:?}, kind={:?})",
            self.kind
        )
    }
}

// ── ML-DSA public key ─────────────────────────────────────────────────────────

/// A validated ML-DSA public key (RFC 9881 / FIPS 204).
///
/// ML-DSA public keys are raw byte strings carried in
/// ``SubjectPublicKeyInfo.subjectPublicKey`` (as BIT STRING value bytes, with
/// the unused-bits prefix stripped).  The expected sizes are:
///
/// - ``ML_DSA_44`` — 1312 bytes
/// - ``ML_DSA_65`` — 1952 bytes
/// - ``ML_DSA_87`` — 2592 bytes
///
/// Use :meth:`from_bytes` with the algorithm OID from
/// ``SubjectPublicKeyInfo.algorithm`` to validate and wrap the raw key bytes:
///
/// ```python,ignore
/// oid = cert.subject_public_key_algorithm
/// raw = cert.subject_public_key
/// pub_key = synta.MlDsaPublicKey.from_bytes(oid, raw)
/// ```
#[pyclass(frozen, name = "MlDsaPublicKey")]
pub struct PyMlDsaPublicKey {
    parameter_set: Py<PyObjectIdentifier>,
    key: Py<PyBytes>,
}

#[pymethods]
impl PyMlDsaPublicKey {
    /// Validate and wrap raw ML-DSA public key bytes.
    ///
    /// ``algorithm_oid`` must be one of ``synta.oids.ML_DSA_44``,
    /// ``synta.oids.ML_DSA_65``, or ``synta.oids.ML_DSA_87`` (or an
    /// equivalent dotted-decimal string).  The byte length of ``data`` is
    /// checked against the size mandated by the OID (1312, 1952, or 2592).
    ///
    /// ``data`` must be the raw key bytes — **not** a DER-encoded OCTET STRING.
    /// When working with a ``SubjectPublicKeyInfo``, pass the value bytes of
    /// the ``subjectPublicKey`` BIT STRING directly (unused-bits byte already
    /// stripped), which :attr:`~synta.Certificate.subject_public_key` provides.
    ///
    /// :param algorithm_oid: :class:`~synta.ObjectIdentifier` or dotted-decimal
    ///     ``str``; must be one of the ``synta.oids.ML_DSA_*`` constants.
    /// :param data: raw public key bytes.
    /// :returns: :class:`MlDsaPublicKey`
    /// :raises ValueError: on unknown OID or wrong byte length.
    ///
    /// ```python,ignore
    /// pub_key = synta.MlDsaPublicKey.from_bytes(synta.oids.ML_DSA_65, raw)
    /// assert len(pub_key) == 1952
    /// ```
    #[staticmethod]
    fn from_bytes(
        py: Python<'_>,
        algorithm_oid: &Bound<'_, PyAny>,
        data: &Bound<'_, PyBytes>,
    ) -> PyResult<Self> {
        use synta_certificate::oids;

        let oid = super::oid_from_pyany(algorithm_oid)?;
        let comps = oid.components();
        let raw = data.as_bytes();

        // Expected sizes per RFC 9881 §3 and FIPS 204.
        let expected_len = if comps == oids::ML_DSA_44 {
            1312usize
        } else if comps == oids::ML_DSA_65 {
            1952
        } else if comps == oids::ML_DSA_87 {
            2592
        } else {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "unknown ML-DSA algorithm OID: {oid}; \
                 expected one of ML_DSA_44, ML_DSA_65, ML_DSA_87",
            )));
        };

        if raw.len() != expected_len {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "ML-DSA public key has wrong length: expected {expected_len} bytes, got {}",
                raw.len()
            )));
        }

        let ps = Py::new(py, PyObjectIdentifier::from_oid(oid))
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
        Ok(Self {
            parameter_set: ps,
            key: PyBytes::new(py, raw).unbind(),
        })
    }

    /// The algorithm :class:`~synta.ObjectIdentifier` identifying the ML-DSA
    /// parameter set (one of ``synta.oids.ML_DSA_44``, ``synta.oids.ML_DSA_65``,
    /// or ``synta.oids.ML_DSA_87``).
    #[getter]
    fn parameter_set<'py>(&self, py: Python<'py>) -> Bound<'py, PyObjectIdentifier> {
        self.parameter_set.bind(py).clone()
    }

    /// The raw public key bytes (1312, 1952, or 2592 octets depending on the
    /// parameter set).
    #[getter]
    fn key<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        self.key.bind(py).clone()
    }

    fn __len__(&self, py: Python<'_>) -> usize {
        self.key.bind(py).as_bytes().len()
    }

    fn __repr__(&self, py: Python<'_>) -> String {
        let oid_str = self.parameter_set.bind(py).borrow().inner.to_string();
        let key_len = self.key.bind(py).as_bytes().len();
        format!("MlDsaPublicKey(parameter_set={oid_str:?}, key=<{key_len} bytes>)")
    }
}

// ── RSA public key ────────────────────────────────────────────────────────────

/// A parsed RSA public key (RFC 8017 / PKCS \#1).
///
/// Holds the two components of an RSA public key:
///
/// - :attr:`modulus` — the RSA modulus ``n`` as big-endian bytes
/// - :attr:`public_exponent` — the public exponent ``e`` as big-endian bytes
///   (commonly ``b'\x01\x00\x01'`` for 65537)
///
/// Use :meth:`from_der` to decode the ``BIT STRING`` payload from a
/// ``SubjectPublicKeyInfo`` structure (after stripping the leading zero byte).
#[pyclass(frozen, name = "RsaPublicKey")]
pub struct PyRsaPublicKey {
    modulus: Py<PyBytes>,
    public_exponent: Py<PyBytes>,
}

#[pymethods]
impl PyRsaPublicKey {
    /// Decode an RSA public key from DER bytes.
    ///
    /// ``data`` must be the DER encoding of ``SEQUENCE { modulus INTEGER,
    /// publicExponent INTEGER }``.  This is the inner payload of the
    /// ``BIT STRING`` in a ``SubjectPublicKeyInfo`` (strip the leading ``0x00``
    /// padding byte before passing here).
    ///
    /// :param data: DER bytes of the ``RSAPublicKey`` SEQUENCE.
    /// :returns: :class:`RsaPublicKey`
    /// :raises ValueError: on malformed DER.
    ///
    /// ```python,ignore
    /// pub_key = synta.RsaPublicKey.from_der(spki_bitstring_value[1:])
    /// print(len(pub_key.modulus))   # e.g. 256 for RSA-2048
    /// ```
    #[staticmethod]
    fn from_der(py: Python<'_>, data: &Bound<'_, PyBytes>) -> PyResult<Self> {
        use synta_certificate::pkcs1_types::RsaPublicKey;
        let raw = data.as_bytes();
        let mut decoder = Decoder::new(raw, Encoding::Der);
        let key: RsaPublicKey = decoder
            .decode()
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
        Ok(Self {
            modulus: PyBytes::new(py, key.modulus.as_bytes()).unbind(),
            public_exponent: PyBytes::new(py, key.public_exponent.as_bytes()).unbind(),
        })
    }

    /// The RSA modulus ``n`` as big-endian bytes.
    #[getter]
    fn modulus<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        self.modulus.bind(py).clone()
    }

    /// The public exponent ``e`` as big-endian bytes (commonly ``b'\x01\x00\x01'``).
    #[getter]
    fn public_exponent<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        self.public_exponent.bind(py).clone()
    }

    fn __repr__(&self, py: Python<'_>) -> String {
        let mod_len = self.modulus.bind(py).as_bytes().len();
        format!("RsaPublicKey(modulus=<{mod_len} bytes>)")
    }
}

// ── RSASSA-PSS algorithm parameters ──────────────────────────────────────────

/// Parsed RSASSA-PSS algorithm parameters (RFC 8017 §A.2.3).
///
/// These appear in the ``parameters`` field of an ``AlgorithmIdentifier``
/// whose ``algorithm`` OID is ``id-RSASSA-PSS`` (``synta.oids.RSASSA_PSS``).
///
/// All fields are ``None`` when the corresponding ASN.1 field is absent (they
/// all have DEFAULT values in the RFC; absent means the RFC default applies).
///
/// - :attr:`hash_algorithm` — OID identifying the hash function (e.g.
///   ``id-sha256``); ``None`` → default SHA-1.
/// - :attr:`mask_gen_algorithm` — OID identifying the mask-generation function
///   (almost always MGF1, ``synta.oids.MGF1``); ``None`` → default MGF1.
/// - :attr:`mask_gen_hash_algorithm` — OID of the hash function *inside* MGF1;
///   decoded from the ``parameters`` of the mask-gen ``AlgorithmIdentifier``.
/// - :attr:`salt_length` — salt length in bytes; ``None`` → default 20.
/// - :attr:`trailer_field` — trailer field integer; ``None`` → default 1 (BC).
#[pyclass(frozen, name = "RsassaPssParams")]
pub struct PyRsassaPssParams {
    hash_algorithm: Option<Py<PyObjectIdentifier>>,
    mask_gen_algorithm: Option<Py<PyObjectIdentifier>>,
    mask_gen_hash_algorithm: Option<Py<PyObjectIdentifier>>,
    salt_length: Option<i64>,
    trailer_field: Option<i64>,
}

#[pymethods]
impl PyRsassaPssParams {
    /// Decode RSASSA-PSS parameters from DER bytes.
    ///
    /// ``data`` must be the DER encoding of the ``RSASSA-PSS-params`` SEQUENCE
    /// from the ``parameters`` field of an ``AlgorithmIdentifier`` whose OID is
    /// ``synta.oids.RSASSA_PSS``.
    ///
    /// :param data: DER bytes of the ``RSASSA-PSS-params`` SEQUENCE.
    /// :returns: :class:`RsassaPssParams`
    /// :raises ValueError: on malformed DER.
    ///
    /// ```python,ignore
    /// # alg_params_der comes from cert.signature_algorithm_parameters_der()
    /// pss = synta.RsassaPssParams.from_der(alg_params_der)
    /// print(pss.hash_algorithm)      # e.g. ObjectIdentifier("2.16.840.1.101.3.4.2.1")
    /// print(pss.salt_length)         # e.g. 32
    /// ```
    #[staticmethod]
    fn from_der(py: Python<'_>, data: &Bound<'_, PyBytes>) -> PyResult<Self> {
        use synta::{Encode, Encoder as SyntaEncoder};
        use synta_certificate::pkcs1_types::RsassaPssParams;
        let raw = data.as_bytes();
        let mut decoder = Decoder::new(raw, Encoding::Der);
        let params: RsassaPssParams<'_> = decoder
            .decode()
            .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;

        let mk_oid = |oid: synta::ObjectIdentifier| -> PyResult<Py<PyObjectIdentifier>> {
            Py::new(py, PyObjectIdentifier::from_oid(oid))
                .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))
        };

        let hash_algorithm = params
            .hash_algorithm
            .map(|a| mk_oid(a.algorithm))
            .transpose()?;

        let (mask_gen_algorithm, mask_gen_hash_algorithm) = match params.mask_gen_algorithm {
            None => (None, None),
            Some(mga) => {
                let mga_oid = mk_oid(mga.algorithm)?;
                // Decode the inner hash AlgorithmIdentifier from mga.parameters.
                let inner_hash_oid = mga.parameters.and_then(|elem| {
                    let mut enc = SyntaEncoder::new(Encoding::Der);
                    elem.encode(&mut enc).ok()?;
                    let bytes = enc.finish().ok()?;
                    let mut dec = Decoder::new(bytes.as_slice(), Encoding::Der);
                    // SAFETY: we decode AlgorithmIdentifier<'_> but bytes is dropped after
                    // the OID is cloned below.  The decode is sound because we clone the OID.
                    let alg_id: synta_certificate::AlgorithmIdentifier<'_> = dec.decode().ok()?;
                    Some(alg_id.algorithm)
                });
                let inner = inner_hash_oid.map(mk_oid).transpose()?;
                (Some(mga_oid), inner)
            }
        };

        let salt_length = params
            .salt_length
            .map(|i| {
                i.as_i64()
                    .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))
            })
            .transpose()?;

        let trailer_field = params
            .trailer_field
            .map(|i| {
                i.as_i64()
                    .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))
            })
            .transpose()?;

        Ok(Self {
            hash_algorithm,
            mask_gen_algorithm,
            mask_gen_hash_algorithm,
            salt_length,
            trailer_field,
        })
    }

    /// OID of the hash function (``None`` → RFC default SHA-1).
    #[getter]
    fn hash_algorithm<'py>(&self, py: Python<'py>) -> Option<Bound<'py, PyObjectIdentifier>> {
        self.hash_algorithm.as_ref().map(|o| o.bind(py).clone())
    }

    /// OID of the mask-generation function (``None`` → RFC default MGF1).
    #[getter]
    fn mask_gen_algorithm<'py>(&self, py: Python<'py>) -> Option<Bound<'py, PyObjectIdentifier>> {
        self.mask_gen_algorithm.as_ref().map(|o| o.bind(py).clone())
    }

    /// OID of the hash function *inside* MGF1 (``None`` if not decoded).
    #[getter]
    fn mask_gen_hash_algorithm<'py>(
        &self,
        py: Python<'py>,
    ) -> Option<Bound<'py, PyObjectIdentifier>> {
        self.mask_gen_hash_algorithm
            .as_ref()
            .map(|o| o.bind(py).clone())
    }

    /// Salt length in bytes (``None`` → RFC default 20).
    #[getter]
    fn salt_length(&self) -> Option<i64> {
        self.salt_length
    }

    /// Trailer field integer (``None`` → RFC default 1).
    #[getter]
    fn trailer_field(&self) -> Option<i64> {
        self.trailer_field
    }

    fn __repr__(&self, py: Python<'_>) -> String {
        let hash = self
            .hash_algorithm
            .as_ref()
            .map(|o| o.bind(py).borrow().inner.to_string())
            .unwrap_or_else(|| "SHA-1 (default)".to_string());
        let salt = self.salt_length.unwrap_or(20);
        format!("RsassaPssParams(hash_algorithm={hash:?}, salt_length={salt})")
    }
}