synta-python-mtc 0.3.1

use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyList};
use synta_mtc::crypto::{MtcProof, MtcSignature};

// ── HashAlgorithm ─────────────────────────────────────────────────────────────

/// Hash algorithm identifier used by Merkle Tree Certificate logs.
///
/// Wraps ``synta_mtc::crypto::HashAlgorithm`` and exposes its properties
/// as Python attributes.
///
/// ```python,ignore
/// import synta.mtc as mtc
///
/// alg = mtc.HashAlgorithm.Sha256
/// print(alg.name())       # "SHA-256"
/// print(alg.output_size()) # 32
/// print(alg.oid())         # "2.16.840.1.101.3.4.2.1"
/// ```
#[pyclass(frozen, name = "HashAlgorithm")]
#[derive(Clone)]
pub struct PyHashAlgorithm {
    pub(super) inner: synta_mtc::crypto::HashAlgorithm,
}

#[pymethods]
impl PyHashAlgorithm {
    /// SHA-256 (32-byte output).
    #[classattr]
    #[allow(non_snake_case)]
    fn Sha256() -> Self {
        PyHashAlgorithm {
            inner: synta_mtc::crypto::HashAlgorithm::Sha256,
        }
    }

    /// SHA-384 (48-byte output).
    #[classattr]
    #[allow(non_snake_case)]
    fn Sha384() -> Self {
        PyHashAlgorithm {
            inner: synta_mtc::crypto::HashAlgorithm::Sha384,
        }
    }

    /// SHA-512 (64-byte output).
    #[classattr]
    #[allow(non_snake_case)]
    fn Sha512() -> Self {
        PyHashAlgorithm {
            inner: synta_mtc::crypto::HashAlgorithm::Sha512,
        }
    }

    /// SHA3-256 (32-byte output).
    #[classattr]
    #[allow(non_snake_case)]
    fn Sha3_256() -> Self {
        PyHashAlgorithm {
            inner: synta_mtc::crypto::HashAlgorithm::Sha3_256,
        }
    }

    /// SHA3-384 (48-byte output).
    #[classattr]
    #[allow(non_snake_case)]
    fn Sha3_384() -> Self {
        PyHashAlgorithm {
            inner: synta_mtc::crypto::HashAlgorithm::Sha3_384,
        }
    }

    /// SHA3-512 (64-byte output).
    #[classattr]
    #[allow(non_snake_case)]
    fn Sha3_512() -> Self {
        PyHashAlgorithm {
            inner: synta_mtc::crypto::HashAlgorithm::Sha3_512,
        }
    }

    /// Return the output size in bytes for this hash algorithm.
    fn output_size(&self) -> usize {
        self.inner.output_size()
    }

    /// Return the canonical name string for this hash algorithm.
    ///
    /// Examples: ``"SHA-256"``, ``"SHA3-384"``.
    fn name(&self) -> &'static str {
        self.inner.name()
    }

    /// Return the dotted-decimal OID string for this hash algorithm.
    ///
    /// Examples: ``"2.16.840.1.101.3.4.2.1"`` for SHA-256.
    fn oid(&self) -> String {
        self.inner.to_oid().to_string()
    }

    fn __repr__(&self) -> String {
        format!(
            "HashAlgorithm.{:?}(name='{}', output_size={})",
            self.inner,
            self.inner.name(),
            self.inner.output_size(),
        )
    }

    fn __str__(&self) -> &'static str {
        self.inner.name()
    }

    fn __eq__(&self, other: &Self) -> bool {
        self.inner == other.inner
    }

    fn __hash__(&self) -> u64 {
        use std::hash::{Hash, Hasher};
        let mut h = std::collections::hash_map::DefaultHasher::new();
        self.inner.hash(&mut h);
        h.finish()
    }
}

// ── MtcSignature ──────────────────────────────────────────────────────────────

/// A single cosignature record embedded in an :class:`MtcProof`.
///
/// Wire format (TLS, big-endian lengths):
///
/// ```text
/// uint8   cosigner_id_len           -- 1-byte length (1..255)
/// opaque  cosigner_id[len]          -- DER-encoded TrustAnchorID (RELATIVE-OID)
/// uint16  signature_value_len
/// opaque  signature_value[len]
/// ```
///
/// ```python,ignore
/// import synta.mtc as mtc
///
/// proof = mtc.MtcProof.decode(raw_bytes)
/// for sig in proof.signatures:
///     print(sig.cosigner_id.hex())
/// ```
#[pyclass(frozen, name = "MtcSignature")]
pub struct PyMtcSignature {
    pub(super) cosigner_id: Vec<u8>,
    pub(super) signature_value: Vec<u8>,
}

fn make_mtc_signature(py: Python<'_>, s: MtcSignature) -> PyResult<Py<PyMtcSignature>> {
    Py::new(
        py,
        PyMtcSignature {
            cosigner_id: s.cosigner_id,
            signature_value: s.signature_value,
        },
    )
}

#[pymethods]
impl PyMtcSignature {
    /// DER-encoded ``TrustAnchorID`` (RELATIVE-OID) identifying the cosigner.
    #[getter]
    fn cosigner_id<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.cosigner_id)
    }

    /// Raw signature bytes produced by the cosigner.
    #[getter]
    fn signature_value<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.signature_value)
    }

    fn __repr__(&self) -> String {
        format!(
            "MtcSignature(cosigner_id_len={}, signature_value_len={})",
            self.cosigner_id.len(),
            self.signature_value.len()
        )
    }

    fn __eq__(&self, other: &Self) -> bool {
        self.cosigner_id == other.cosigner_id && self.signature_value == other.signature_value
    }
}

// ── MtcProof ──────────────────────────────────────────────────────────────────

/// TLS-encoded ``MTCProof`` embedded as the X.509 ``signatureValue`` field.
///
/// Per draft-ietf-plants-merkle-tree-certs-04 §4.3, the proof contains:
/// - ``extensions`` — raw ``MerkleTreeCertEntryExtensions`` bytes
/// - ``start`` / ``end`` — uint48 subtree range (first and last+1 leaf indices)
/// - ``inclusion_proof`` — concatenated sibling hashes from the inclusion path
/// - ``signatures`` — list of :class:`MtcSignature` cosignature records
///
/// Use :meth:`decode` to parse from TLS wire bytes and :meth:`encode` to
/// re-serialise.
///
/// ```python,ignore
/// import synta.mtc as mtc
///
/// proof = mtc.MtcProof.decode(raw_bytes)
/// print(proof.start, proof.end)
/// print(len(proof.signatures), "cosignatures")
/// re_encoded = proof.encode()
/// assert re_encoded == raw_bytes
/// ```
#[pyclass(frozen, name = "MtcProof")]
pub struct PyMtcProof {
    pub(super) extensions: Vec<u8>,
    pub(super) start: u64,
    pub(super) end: u64,
    pub(super) inclusion_proof: Vec<u8>,
    pub(super) signatures: Vec<Py<PyMtcSignature>>,
}

#[pymethods]
impl PyMtcProof {
    /// Decode an ``MtcProof`` from its TLS wire format.
    ///
    /// :raises ValueError: if *data* is truncated, has trailing bytes, or
    ///     contains cosigner IDs that are not in canonical order per spec §6.1.
    #[staticmethod]
    fn decode(py: Python<'_>, data: &[u8]) -> PyResult<Self> {
        let proof = MtcProof::decode(data).map_err(|e| PyValueError::new_err(e.to_string()))?;
        let signatures = proof
            .signatures
            .into_iter()
            .map(|s| make_mtc_signature(py, s))
            .collect::<PyResult<Vec<_>>>()?;
        Ok(PyMtcProof {
            extensions: proof.extensions,
            start: proof.start,
            end: proof.end,
            inclusion_proof: proof.inclusion_proof,
            signatures,
        })
    }

    /// Encode this ``MtcProof`` to its TLS wire format.
    ///
    /// :raises ValueError: if any length field overflows its TLS encoding limit,
    ///     or if the cosigner IDs are not in canonical order per spec §6.1.
    fn encode<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
        let signatures: Vec<MtcSignature> = self
            .signatures
            .iter()
            .map(|s| {
                let s = s.get();
                MtcSignature {
                    cosigner_id: s.cosigner_id.clone(),
                    signature_value: s.signature_value.clone(),
                }
            })
            .collect();
        let proof = MtcProof {
            extensions: self.extensions.clone(),
            start: self.start,
            end: self.end,
            inclusion_proof: self.inclusion_proof.clone(),
            signatures,
        };
        let bytes = proof
            .encode()
            .map_err(|e| PyValueError::new_err(e.to_string()))?;
        Ok(PyBytes::new(py, &bytes))
    }

    /// Raw ``MerkleTreeCertEntryExtensions`` bytes (empty when no extensions).
    #[getter]
    fn extensions<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.extensions)
    }

    /// First leaf index of the subtree (inclusive, uint48).
    #[getter]
    fn start(&self) -> u64 {
        self.start
    }

    /// Last leaf index of the subtree (exclusive, uint48).
    #[getter]
    fn end(&self) -> u64 {
        self.end
    }

    /// Concatenated sibling hashes from the inclusion proof path.
    #[getter]
    fn inclusion_proof<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.inclusion_proof)
    }

    /// List of :class:`MtcSignature` cosignature records, in canonical order.
    #[getter]
    fn signatures<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyList>> {
        let elems: Vec<Bound<'_, PyMtcSignature>> = self
            .signatures
            .iter()
            .map(|x| x.clone_ref(py).into_bound(py))
            .collect();
        PyList::new(py, elems)
    }

    fn __repr__(&self) -> String {
        format!(
            "MtcProof(start={}, end={}, inclusion_proof_len={}, signatures={})",
            self.start,
            self.end,
            self.inclusion_proof.len(),
            self.signatures.len()
        )
    }

    fn __eq__(&self, other: &Self) -> bool {
        self.extensions == other.extensions
            && self.start == other.start
            && self.end == other.end
            && self.inclusion_proof == other.inclusion_proof
            && self.signatures.len() == other.signatures.len()
    }
}