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;
use synta_python_common::SyntaErr;

// ── ValidationPolicy ─────────────────────────────────────────────────────────

/// Validation policy for Merkle Tree Certificates.
///
/// Controls which hash algorithms are allowed, cosignature quorum requirements,
/// and other validation constraints.  Use :class:`CertificateValidator` to apply
/// a policy during certificate validation.
///
/// The default policy allows only SHA-256 and requires at least one cosignature.
/// Use :meth:`permissive` for testing (all hash algorithms, no cosignature
/// requirement).
///
/// ```python,ignore
/// import synta.mtc as mtc
///
/// policy = mtc.ValidationPolicy()
/// policy.allow_hash_algorithm(mtc.HashAlgorithm.Sha384)
/// policy.set_min_cosignatures(2)
/// ```
#[pyclass(name = "ValidationPolicy")]
pub struct PyValidationPolicy {
    /// Allowed hash algorithm OID strings (dotted-decimal).
    allowed_hash_algorithm_oids: Vec<String>,
    /// Minimum number of valid cosignatures required.
    min_cosignatures: usize,
    /// Maximum number of cosignatures allowed (DoS protection).
    max_cosignatures: usize,
    /// Whether duplicate cosigners are allowed.
    allow_duplicate_cosigners: bool,
    /// Whether to require valid timestamps.
    require_valid_timestamps: bool,
    /// Whether to check certificate expiration.
    check_expiration: bool,
    /// Maximum certificate age in seconds, or ``None``.
    max_certificate_age: Option<u64>,
}

impl PyValidationPolicy {
    /// Build a ``synta_mtc::config::ValidationPolicy`` from this Python wrapper.
    fn to_rust(&self) -> synta_mtc::config::ValidationPolicy {
        use std::str::FromStr;
        use synta_mtc::config::ValidationPolicy;

        let mut policy = ValidationPolicy::new();
        policy.allowed_hash_algorithms.clear();
        for oid_str in &self.allowed_hash_algorithm_oids {
            if let Ok(oid) = synta::ObjectIdentifier::from_str(oid_str) {
                policy.allowed_hash_algorithms.push(oid);
            }
        }
        policy.cosignature_policy.min_cosignatures = self.min_cosignatures;
        policy.cosignature_policy.max_cosignatures = self.max_cosignatures;
        policy.cosignature_policy.allow_duplicate_cosigners = self.allow_duplicate_cosigners;
        policy.require_valid_timestamps = self.require_valid_timestamps;
        policy.check_expiration = self.check_expiration;
        policy.max_certificate_age = self.max_certificate_age;
        policy
    }
}

#[pymethods]
impl PyValidationPolicy {
    /// Create a new ``ValidationPolicy`` with default settings.
    ///
    /// Default: SHA-256 allowed, at least 1 cosignature required, timestamps
    /// and expiration checked.
    #[new]
    fn new() -> Self {
        use synta_mtc::config::ValidationPolicy;
        let default = ValidationPolicy::default();
        PyValidationPolicy {
            allowed_hash_algorithm_oids: default
                .allowed_hash_algorithms
                .iter()
                .map(|o| o.to_string())
                .collect(),
            min_cosignatures: default.cosignature_policy.min_cosignatures,
            max_cosignatures: default.cosignature_policy.max_cosignatures,
            allow_duplicate_cosigners: default.cosignature_policy.allow_duplicate_cosigners,
            require_valid_timestamps: default.require_valid_timestamps,
            check_expiration: default.check_expiration,
            max_certificate_age: default.max_certificate_age,
        }
    }

    /// Create a permissive ``ValidationPolicy`` for testing.
    ///
    /// Allows all supported hash algorithms, requires zero cosignatures, and
    /// disables timestamp and expiration checks.
    #[staticmethod]
    fn permissive() -> Self {
        use synta_mtc::config::ValidationPolicy;
        let p = ValidationPolicy::permissive();
        PyValidationPolicy {
            allowed_hash_algorithm_oids: p
                .allowed_hash_algorithms
                .iter()
                .map(|o| o.to_string())
                .collect(),
            min_cosignatures: p.cosignature_policy.min_cosignatures,
            max_cosignatures: p.cosignature_policy.max_cosignatures,
            allow_duplicate_cosigners: p.cosignature_policy.allow_duplicate_cosigners,
            require_valid_timestamps: p.require_valid_timestamps,
            check_expiration: p.check_expiration,
            max_certificate_age: p.max_certificate_age,
        }
    }

    /// Allow an additional hash algorithm.
    ///
    /// :param algo: A :class:`HashAlgorithm` instance (e.g. ``HashAlgorithm.Sha384``).
    fn allow_hash_algorithm(&mut self, algo: &PyHashAlgorithm) {
        let oid_str = algo.inner.to_oid().to_string();
        if !self.allowed_hash_algorithm_oids.contains(&oid_str) {
            self.allowed_hash_algorithm_oids.push(oid_str);
        }
    }

    /// Set the minimum number of cosignatures required for validation.
    ///
    /// :param n: Minimum count (0 = no cosignatures required).
    fn set_min_cosignatures(&mut self, n: usize) {
        self.min_cosignatures = n;
    }

    /// Set the maximum number of cosignatures accepted (DoS protection).
    ///
    /// :param n: Maximum count.
    fn set_max_cosignatures(&mut self, n: usize) {
        self.max_cosignatures = n;
    }

    /// Allow or forbid duplicate cosigner IDs in the same certificate.
    ///
    /// :param allow: When ``True``, a single cosigner may contribute multiple
    ///     entries toward the quorum (insecure — leave ``False`` in production).
    fn set_allow_duplicate_cosigners(&mut self, allow: bool) {
        self.allow_duplicate_cosigners = allow;
    }

    /// Enable or disable timestamp validity checking.
    ///
    /// :param require: When ``True``, certificates with invalid timestamps are rejected.
    fn set_require_valid_timestamps(&mut self, require: bool) {
        self.require_valid_timestamps = require;
    }

    /// Enable or disable certificate expiration checking.
    ///
    /// :param check: When ``True``, expired certificates are rejected.
    fn set_check_expiration(&mut self, check: bool) {
        self.check_expiration = check;
    }

    /// Allowed hash algorithm OID strings (dotted-decimal).
    #[getter]
    fn allowed_hash_algorithm_oids(&self) -> Vec<String> {
        self.allowed_hash_algorithm_oids.clone()
    }

    /// Minimum number of valid cosignatures required.
    #[getter]
    fn min_cosignatures(&self) -> usize {
        self.min_cosignatures
    }

    /// Maximum number of cosignatures accepted (DoS protection).
    #[getter]
    fn max_cosignatures(&self) -> usize {
        self.max_cosignatures
    }

    /// Whether duplicate cosigners are allowed.
    #[getter]
    fn allow_duplicate_cosigners(&self) -> bool {
        self.allow_duplicate_cosigners
    }

    /// Whether timestamp validity is checked.
    #[getter]
    fn require_valid_timestamps(&self) -> bool {
        self.require_valid_timestamps
    }

    /// Whether certificate expiration is checked.
    #[getter]
    fn check_expiration(&self) -> bool {
        self.check_expiration
    }

    fn __repr__(&self) -> String {
        format!(
            "ValidationPolicy(min_cosignatures={}, algorithms={})",
            self.min_cosignatures,
            self.allowed_hash_algorithm_oids.len()
        )
    }
}

// ── TrustAnchor ───────────────────────────────────────────────────────────────

/// A trust anchor for a Merkle Tree Certificate log.
///
/// Identifies a log operator by their ``LogID`` (hash algorithm + public key),
/// provides a human-readable name for diagnostics, tracks whether the anchor is
/// active, and carries an optional revocation list (spec §7.5).
///
/// Construct with :meth:`from_log_id_der` and add to a
/// :class:`CertificateValidator` via :meth:`CertificateValidator.add_trust_anchor`.
///
/// ```python,ignore
/// import synta.mtc as mtc
///
/// anchor = mtc.TrustAnchor.from_log_id_der(log_id_der, "My Log")
/// anchor.revoke_range(1, 100)
///
/// validator = mtc.CertificateValidator()
/// validator.add_trust_anchor(anchor)
/// ```
#[pyclass(name = "TrustAnchor")]
pub struct PyTrustAnchor {
    /// Owned DER bytes for the LogID (keeps the decoded LogID alive).
    pub(super) log_id_der: Box<[u8]>,
    /// Human-readable name for diagnostics.
    pub(super) name: String,
    /// Whether this anchor is currently active.
    pub(super) active: bool,
    /// Revocation ranges for this anchor.
    pub(super) revoked_ranges: synta_mtc::revocation::RevokedRanges,
}

#[pymethods]
impl PyTrustAnchor {
    /// Construct a ``TrustAnchor`` from a DER-encoded ``LogID`` and a display name.
    ///
    /// :param log_id_der: DER-encoded ``LogID`` SEQUENCE bytes.
    /// :param name: Human-readable name for diagnostics.
    /// :raises ValueError: if *log_id_der* cannot be decoded as a ``LogID``.
    #[staticmethod]
    fn from_log_id_der(log_id_der: &[u8], name: String) -> PyResult<Self> {
        // Validate by decoding.
        let validated = {
            let mut dec = synta::Decoder::new(log_id_der, synta::Encoding::Der);
            synta_mtc::types::LogID::decode(&mut dec).map_err(SyntaErr)?
        };
        // Re-encode to get a fully owned, canonical copy.
        let mut enc = synta::Encoder::new(synta::Encoding::Der);
        validated.encode(&mut enc).map_err(SyntaErr)?;
        let owned: Box<[u8]> = enc.finish().map_err(SyntaErr)?.into_boxed_slice();
        Ok(PyTrustAnchor {
            log_id_der: owned,
            name,
            active: true,
            revoked_ranges: synta_mtc::revocation::RevokedRanges::new(),
        })
    }

    /// Return the DER-encoded ``LogID`` for this trust anchor.
    #[getter]
    fn log_id_der<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
        PyBytes::new(py, &self.log_id_der)
    }

    /// Human-readable name for this trust anchor (for diagnostics).
    #[getter]
    fn name(&self) -> &str {
        &self.name
    }

    /// Whether this trust anchor is currently active.
    ///
    /// Inactive anchors are ignored during validation.
    #[getter]
    fn active(&self) -> bool {
        self.active
    }

    /// Activate this trust anchor.
    fn activate(&mut self) {
        self.active = true;
    }

    /// Deactivate this trust anchor.
    ///
    /// Deactivated anchors are skipped by the validator.
    fn deactivate(&mut self) {
        self.active = false;
    }

    /// Return ``True`` if *entry_index* falls within a revoked range (spec §7.5).
    ///
    /// :param entry_index: Leaf index to check.
    fn is_revoked(&self, entry_index: u64) -> bool {
        self.revoked_ranges.is_revoked(entry_index)
    }

    /// Add an inclusive range ``[start, end]`` of revoked entry indices (spec §7.5).
    ///
    /// Overlapping and adjacent ranges are automatically merged.
    ///
    /// :param start: Start of the revoked range (inclusive).
    /// :param end: End of the revoked range (inclusive).
    /// :raises ValueError: if *start* > *end*.
    fn revoke_range(&mut self, start: u64, end: u64) -> PyResult<()> {
        self.revoked_ranges
            .add(start, end)
            .map_err(|e| PyValueError::new_err(e.to_string()))
    }

    fn __repr__(&self) -> String {
        format!("TrustAnchor(name='{}', active={})", self.name, self.active)
    }
}

// ── CertificateValidator ──────────────────────────────────────────────────────

/// Validates Merkle Tree Certificates against trusted checkpoints.
///
/// Combines trust anchor lookup, inclusion proof verification, subtree
/// consistency checking, and policy enforcement.
///
/// Construction: create a validator, add trust anchors, then call
/// :meth:`validate_standalone` or :meth:`validate_landmark`.
///
/// ```python,ignore
/// import synta.mtc as mtc
///
/// anchor = mtc.TrustAnchor.from_log_id_der(log_id_der, "My Log")
///
/// policy = mtc.ValidationPolicy()
/// policy.allow_hash_algorithm(mtc.HashAlgorithm.Sha256)
///
/// validator = mtc.CertificateValidator(policy)
/// validator.add_trust_anchor(anchor)
///
/// try:
///     validator.validate_standalone(standalone_cert_der)
///     print("certificate is valid")
/// except ValueError as e:
///     print("invalid:", e)
/// ```
#[pyclass(name = "CertificateValidator")]
pub struct PyCertificateValidator {
    /// Owned DER bytes for each trust anchor's LogID.
    /// Each `Box<[u8]>` has a stable heap address that the decoded
    /// `LogID<'static>` borrows from.
    anchor_log_id_bytes: Vec<Box<[u8]>>,
    /// Trust anchor metadata (name, active flag, revoked ranges), parallel to
    /// `anchor_log_id_bytes`.
    anchor_meta: Vec<(String, bool, synta_mtc::revocation::RevokedRanges)>,
    /// Validation policy stored as a Python wrapper (cheap to clone).
    policy: PyValidationPolicy,
}

impl PyCertificateValidator {
    /// Build a ``CertificateValidator<'static>`` whose trust anchors borrow from
    /// the stable ``Box<[u8]>`` heap addresses in ``self.anchor_log_id_bytes``.
    ///
    /// SAFETY: the returned validator must not outlive `self` because the
    /// `LogID<'static>` values borrow from boxes owned by `self`.
    fn build_rust_validator(
        &self,
    ) -> PyResult<synta_mtc::validator::CertificateValidator<'static>> {
        use synta_mtc::config::{TrustAnchor, ValidatorConfig};
        use synta_mtc::types::LogID;

        let policy = self.policy.to_rust();
        let mut config = ValidatorConfig::with_policy(policy);

        for (der_box, (name, active, revoked)) in
            self.anchor_log_id_bytes.iter().zip(self.anchor_meta.iter())
        {
            // SAFETY: the Box heap address is stable for the life of `self`.
            let static_bytes: &'static [u8] = unsafe { &*(der_box.as_ref() as *const [u8]) };
            let mut dec = synta::Decoder::new(static_bytes, synta::Encoding::Der);
            let log_id_decoded = LogID::decode(&mut dec).map_err(SyntaErr)?;
            // SAFETY: `log_id_decoded` borrows `static_bytes`, which is backed by
            // the Box heap address that outlives the validator we build here
            // (the validator is used synchronously and dropped before any Box moves).
            let log_id: LogID<'static> = unsafe { std::mem::transmute(log_id_decoded) };
            let mut anchor = TrustAnchor::new(log_id, name.clone());
            anchor.active = *active;
            anchor.revoked_ranges = revoked.clone();
            config.add_trust_anchor(anchor);
        }

        Ok(synta_mtc::validator::CertificateValidator::new(config))
    }
}

#[pymethods]
impl PyCertificateValidator {
    /// Create a new ``CertificateValidator`` with an optional policy.
    ///
    /// If *policy* is omitted the default policy is used (SHA-256 only,
    /// at least one cosignature required).
    ///
    /// :param policy: A :class:`ValidationPolicy` instance, or ``None`` for defaults.
    #[new]
    #[pyo3(signature = (policy=None))]
    fn new(policy: Option<&PyValidationPolicy>) -> Self {
        let policy = match policy {
            Some(p) => PyValidationPolicy {
                allowed_hash_algorithm_oids: p.allowed_hash_algorithm_oids.clone(),
                min_cosignatures: p.min_cosignatures,
                max_cosignatures: p.max_cosignatures,
                allow_duplicate_cosigners: p.allow_duplicate_cosigners,
                require_valid_timestamps: p.require_valid_timestamps,
                check_expiration: p.check_expiration,
                max_certificate_age: p.max_certificate_age,
            },
            None => PyValidationPolicy::new(),
        };
        PyCertificateValidator {
            anchor_log_id_bytes: Vec::new(),
            anchor_meta: Vec::new(),
            policy,
        }
    }

    /// Add a trust anchor to the validator.
    ///
    /// The anchor's ``LogID`` DER bytes are copied into stable owned storage
    /// so that the anchor can be used across calls without holding a reference
    /// to the original :class:`TrustAnchor` object.
    ///
    /// :param anchor: A :class:`TrustAnchor` instance.
    fn add_trust_anchor(&mut self, anchor: &PyTrustAnchor) {
        self.anchor_log_id_bytes.push(anchor.log_id_der.clone());
        self.anchor_meta.push((
            anchor.name.clone(),
            anchor.active,
            anchor.revoked_ranges.clone(),
        ));
    }

    /// Return the number of trust anchors currently registered.
    fn trust_anchor_count(&self) -> usize {
        self.anchor_log_id_bytes.len()
    }

    /// Validate a DER-encoded ``StandaloneCertificate``.
    ///
    /// Runs the full validation pipeline (spec §4.3.3):
    ///
    /// 1. Cosignature quorum check.
    /// 2. Find a trusted cosignature.
    /// 3. Resolve and policy-check the hash algorithm.
    /// 4. Verify the leaf index falls within the cosigner's subtree.
    /// 5. Check the leaf index against the trust anchor's revocation list (§7.5).
    /// 6. Compute the leaf hash from the embedded TBS certificate.
    /// 7. Verify the embedded inclusion proof against the trusted checkpoint root.
    ///
    /// :param cert_der: DER-encoded ``StandaloneCertificate`` bytes.
    /// :raises ValueError: if validation fails for any reason.
    fn validate_standalone(&self, cert_der: &[u8]) -> PyResult<()> {
        use synta_mtc::types::StandaloneCertificate;

        let mut dec = synta::Decoder::new(cert_der, synta::Encoding::Der);
        let cert = StandaloneCertificate::decode(&mut dec).map_err(SyntaErr)?;
        let validator = self.build_rust_validator()?;
        validator
            .validate_standalone(&cert)
            .map_err(|e| PyValueError::new_err(e.to_string()))
    }

    /// Validate a DER-encoded ``LandmarkCertificate`` against a trusted checkpoint.
    ///
    /// Landmark certificates carry no embedded cosignatures. You must supply
    /// the trusted :class:`Checkpoint` DER bytes (obtained from a separately
    /// validated subtree proof for the log referenced by the certificate's
    /// ``landmark_id``).
    ///
    /// Validation pipeline (spec §4.3.3):
    ///
    /// 1. Verify that ``landmark_id.log_id`` matches an active trust anchor.
    /// 2. Resolve and policy-check the hash algorithm.
    /// 3. Extract the leaf index from the serial number (§6.1).
    /// 4. Check the leaf index against the trust anchor's revocation list (§7.5).
    /// 5. Compute the leaf hash from the embedded TBS certificate.
    /// 6. Verify the inclusion proof against the supplied checkpoint's root.
    ///
    /// :param cert_der: DER-encoded ``LandmarkCertificate`` bytes.
    /// :param checkpoint_der: DER-encoded ``Checkpoint`` bytes (must be trusted).
    /// :raises ValueError: if validation fails for any reason.
    fn validate_landmark(&self, cert_der: &[u8], checkpoint_der: &[u8]) -> PyResult<()> {
        use synta_mtc::types::{Checkpoint, LandmarkCertificate};

        let cert = {
            let mut dec = synta::Decoder::new(cert_der, synta::Encoding::Der);
            LandmarkCertificate::decode(&mut dec).map_err(SyntaErr)?
        };
        let checkpoint = {
            let mut dec = synta::Decoder::new(checkpoint_der, synta::Encoding::Der);
            Checkpoint::decode(&mut dec).map_err(SyntaErr)?
        };
        let validator = self.build_rust_validator()?;
        validator
            .validate_landmark(&cert, &checkpoint)
            .map_err(|e| PyValueError::new_err(e.to_string()))
    }

    fn __repr__(&self) -> String {
        format!(
            "CertificateValidator(trust_anchors={}, min_cosignatures={})",
            self.anchor_log_id_bytes.len(),
            self.policy.min_cosignatures
        )
    }
}