xdoc-rs 0.1.1

use crate::core::{Attribute, Document, ErrorKind, NodeId, NodeKind, QName, XmlError, XmlResult};

use super::canonicalization::canonicalize_node_excluding;
use super::xmldsig::{
    element_children, find_signature, optional_attribute, required_child, XMLDSIG_NAMESPACE_URI,
};
use super::{
    decode_standard_base64, digest_bytes, encode_standard_base64, CanonicalizationConfig,
    DigestAlgorithm, TimestampAuthorityClient, TimestampRequest, TimestampToken,
    XADES_NAMESPACE_URI,
};

/// Configuration for XAdES archive timestamps.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct XadesArchiveConfig {
    digest_algorithm: DigestAlgorithm,
    canonicalization: CanonicalizationConfig,
}

impl XadesArchiveConfig {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_digest_algorithm(mut self, digest_algorithm: DigestAlgorithm) -> Self {
        self.digest_algorithm = digest_algorithm;
        self
    }

    pub fn with_canonicalization(mut self, canonicalization: CanonicalizationConfig) -> Self {
        self.canonicalization = canonicalization;
        self
    }

    pub fn digest_algorithm(&self) -> DigestAlgorithm {
        self.digest_algorithm
    }

    pub fn canonicalization(&self) -> &CanonicalizationConfig {
        &self.canonicalization
    }
}

impl Default for XadesArchiveConfig {
    fn default() -> Self {
        Self {
            digest_algorithm: DigestAlgorithm::Sha256,
            canonicalization: CanonicalizationConfig::default(),
        }
    }
}

/// Report for XAdES archive timestamp preservation state.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct XadesArchiveReport {
    pub signature_timestamp_present: bool,
    pub validation_data_present: bool,
    pub archive_timestamp_count: usize,
    pub structure_valid: bool,
    pub token_valid_count: usize,
    pub latest_token_valid: bool,
    pub preservation_ready: bool,
    pub message_imprint: Vec<u8>,
}

/// Adds a new `xades:ArchiveTimeStamp` unsigned property.
///
/// The timestamp request is computed over the signature node after excluding
/// existing archive timestamps. This keeps re-timestamping deterministic while
/// preserving the original signed XMLDSig/XAdES material.
pub fn add_xades_archive_timestamp<C>(
    document: &Document,
    client: &C,
    config: &XadesArchiveConfig,
) -> XmlResult<Document>
where
    C: TimestampAuthorityClient + ?Sized,
{
    config.digest_algorithm.ensure_allowed_for_generation()?;

    let mut archived = document.clone();
    let signature = find_signature(&archived)?;
    let preservation_inputs = preservation_inputs_present(&archived, signature)?;
    if !preservation_inputs.signature_timestamp_present {
        return Err(XmlError::new(
            ErrorKind::Signature,
            "XAdES archive timestamp requires an existing SignatureTimeStamp",
        ));
    }
    if !preservation_inputs.validation_data_present {
        return Err(XmlError::new(
            ErrorKind::Signature,
            "XAdES archive timestamp requires existing validation data",
        ));
    }

    let message_imprint = archive_message_imprint(&archived, signature, config)?;
    let request = TimestampRequest {
        digest_algorithm: config.digest_algorithm,
        message_imprint,
    };
    let token = client.timestamp(&request)?;
    let qualifying_properties = find_qualifying_properties(&archived, signature)?;
    let unsigned_signature_properties =
        ensure_unsigned_signature_properties(&mut archived, qualifying_properties)?;

    let archive_timestamp = archived.add_element(
        unsigned_signature_properties,
        QName::qualified("xades", "ArchiveTimeStamp", XADES_NAMESPACE_URI)?,
    )?;
    let canonicalization = archived.add_element(
        archive_timestamp,
        QName::qualified("ds", "CanonicalizationMethod", XMLDSIG_NAMESPACE_URI)?,
    )?;
    archived.add_attribute(
        canonicalization,
        Attribute::new(
            QName::new("Algorithm")?,
            config.canonicalization.algorithm().uri(),
        ),
    )?;
    let encoded = archived.add_element(
        archive_timestamp,
        QName::qualified("xades", "EncapsulatedTimeStamp", XADES_NAMESPACE_URI)?,
    )?;
    archived.add_text(encoded, encode_standard_base64(token.encoded))?;

    Ok(archived)
}

/// Verifies all existing `xades:ArchiveTimeStamp` tokens against the preserved signature state.
pub fn verify_xades_archive_timestamps<C>(
    document: &Document,
    client: &C,
    config: &XadesArchiveConfig,
) -> XmlResult<XadesArchiveReport>
where
    C: TimestampAuthorityClient + ?Sized,
{
    config.digest_algorithm.ensure_allowed_for_generation()?;

    let signature = find_signature(document)?;
    let preservation_inputs = preservation_inputs_present(document, signature)?;
    let archive_timestamps = find_archive_timestamps(document, signature)?;
    let message_imprint = archive_message_imprint(document, signature, config)?;
    let request = TimestampRequest {
        digest_algorithm: config.digest_algorithm,
        message_imprint: message_imprint.clone(),
    };

    let mut structure_valid = true;
    let mut token_valid_count = 0;
    let mut latest_token_valid = false;
    for timestamp in &archive_timestamps {
        let canonicalization_method =
            required_child(document, *timestamp, "CanonicalizationMethod")?;
        let algorithm = optional_attribute(document, canonicalization_method, "Algorithm")?
            .ok_or_else(|| {
                XmlError::new(
                    ErrorKind::Signature,
                    "XAdES ArchiveTimeStamp CanonicalizationMethod requires Algorithm",
                )
            })?;
        let timestamp_structure_valid = algorithm == config.canonicalization.algorithm().uri();
        structure_valid = structure_valid && timestamp_structure_valid;
        let token_text = child_text(document, *timestamp, "EncapsulatedTimeStamp")?;
        let token = TimestampToken::new(decode_standard_base64(&token_text)?);
        let token_valid = timestamp_structure_valid && client.verify(&request, &token)?;
        if token_valid {
            token_valid_count += 1;
        }
        latest_token_valid = token_valid;
    }

    let preservation_ready = preservation_inputs.signature_timestamp_present
        && preservation_inputs.validation_data_present
        && !archive_timestamps.is_empty()
        && structure_valid
        && token_valid_count == archive_timestamps.len();

    Ok(XadesArchiveReport {
        signature_timestamp_present: preservation_inputs.signature_timestamp_present,
        validation_data_present: preservation_inputs.validation_data_present,
        archive_timestamp_count: archive_timestamps.len(),
        structure_valid,
        token_valid_count,
        latest_token_valid,
        preservation_ready,
        message_imprint,
    })
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct PreservationInputs {
    signature_timestamp_present: bool,
    validation_data_present: bool,
}

fn preservation_inputs_present(
    document: &Document,
    signature: NodeId,
) -> XmlResult<PreservationInputs> {
    let qualifying_properties = find_qualifying_properties(document, signature)?;
    let Some(unsigned_properties) =
        optional_xades_child(document, qualifying_properties, "UnsignedProperties")?
    else {
        return Ok(PreservationInputs {
            signature_timestamp_present: false,
            validation_data_present: false,
        });
    };
    let Some(unsigned_signature_properties) =
        optional_xades_child(document, unsigned_properties, "UnsignedSignatureProperties")?
    else {
        return Ok(PreservationInputs {
            signature_timestamp_present: false,
            validation_data_present: false,
        });
    };
    let signature_timestamp_present = optional_xades_child(
        document,
        unsigned_signature_properties,
        "SignatureTimeStamp",
    )?
    .is_some();
    let certificate_values_present =
        optional_xades_child(document, unsigned_signature_properties, "CertificateValues")?
            .is_some();
    let revocation_values_present =
        optional_xades_child(document, unsigned_signature_properties, "RevocationValues")?
            .is_some();

    Ok(PreservationInputs {
        signature_timestamp_present,
        validation_data_present: certificate_values_present && revocation_values_present,
    })
}

fn archive_message_imprint(
    document: &Document,
    signature: NodeId,
    config: &XadesArchiveConfig,
) -> XmlResult<Vec<u8>> {
    let archive_timestamps = find_archive_timestamps(document, signature)?;
    let canonicalized = canonicalize_node_excluding(
        document,
        signature,
        &archive_timestamps,
        &config.canonicalization,
    )?;
    digest_bytes(config.digest_algorithm, canonicalized)
}

fn find_archive_timestamps(document: &Document, signature: NodeId) -> XmlResult<Vec<NodeId>> {
    let qualifying_properties = find_qualifying_properties(document, signature)?;
    let Some(unsigned_properties) =
        optional_xades_child(document, qualifying_properties, "UnsignedProperties")?
    else {
        return Ok(Vec::new());
    };
    let Some(unsigned_signature_properties) =
        optional_xades_child(document, unsigned_properties, "UnsignedSignatureProperties")?
    else {
        return Ok(Vec::new());
    };
    Ok(element_children(document, unsigned_signature_properties)?
        .into_iter()
        .filter(|child| is_xades_element(document, *child, "ArchiveTimeStamp"))
        .collect())
}

fn find_qualifying_properties(document: &Document, signature: NodeId) -> XmlResult<NodeId> {
    let object = required_child(document, signature, "Object")?;
    element_children(document, object)?
        .into_iter()
        .find(|child| is_xades_element(document, *child, "QualifyingProperties"))
        .ok_or_else(|| {
            XmlError::new(
                ErrorKind::Signature,
                "missing required XAdES QualifyingProperties",
            )
        })
}

fn ensure_unsigned_signature_properties(
    document: &mut Document,
    qualifying_properties: NodeId,
) -> XmlResult<NodeId> {
    let unsigned_properties =
        match optional_xades_child(document, qualifying_properties, "UnsignedProperties")? {
            Some(node) => node,
            None => document.add_element(
                qualifying_properties,
                QName::qualified("xades", "UnsignedProperties", XADES_NAMESPACE_URI)?,
            )?,
        };

    match optional_xades_child(document, unsigned_properties, "UnsignedSignatureProperties")? {
        Some(node) => Ok(node),
        None => document.add_element(
            unsigned_properties,
            QName::qualified("xades", "UnsignedSignatureProperties", XADES_NAMESPACE_URI)?,
        ),
    }
}

fn optional_xades_child(
    document: &Document,
    parent: NodeId,
    local: &str,
) -> XmlResult<Option<NodeId>> {
    Ok(element_children(document, parent)?
        .into_iter()
        .find(|child| is_xades_element(document, *child, local)))
}

fn child_text(document: &Document, parent: NodeId, local: &str) -> XmlResult<String> {
    let child = required_child(document, parent, local)?;
    let mut text = String::new();
    for node in document.children(child)? {
        if let NodeKind::Text(value) = document.node(*node)?.kind() {
            text.push_str(value);
        }
    }
    if text.is_empty() {
        return Err(XmlError::new(
            ErrorKind::Signature,
            format!("XAdES ArchiveTimeStamp child `{local}` must contain text"),
        ));
    }
    Ok(text)
}

fn is_xades_element(document: &Document, node: NodeId, local: &str) -> bool {
    matches!(
        document.node(node).map(|node| node.kind()),
        Ok(NodeKind::Element(element))
            if element.name().namespace_uri().map(|uri| uri.as_str()) == Some(XADES_NAMESPACE_URI)
                && element.name().local() == local
    )
}

#[cfg(test)]
mod tests {
    use crate::parser::parse_str;
    use crate::signature::{
        add_signature_timestamp, add_xades_validation_data, verify_xades_bes_enveloped,
        DeterministicSigningProvider, DeterministicTimestampAuthority,
        StaticValidationDataProvider, XadesConfig, XadesTimestampConfig, XadesValidationDataConfig,
    };
    use crate::writer::to_string_compact;

    use super::*;

    fn signing_provider() -> DeterministicSigningProvider {
        DeterministicSigningProvider::new(b"test-cert".to_vec(), b"test-secret".to_vec())
    }

    fn timestamp_authority() -> DeterministicTimestampAuthority {
        DeterministicTimestampAuthority::new(b"timestamp-secret")
    }

    fn signed_document() -> XmlResult<Document> {
        let document = parse_str(r#"<Root Id="doc-1"><Item>value</Item></Root>"#)?;
        crate::signature::sign_xades_bes_enveloped(
            &document,
            &signing_provider(),
            &XadesConfig::new().with_signing_time("2026-06-11T12:00:00Z"),
        )
    }

    fn preserved_document() -> XmlResult<Document> {
        let timestamped = add_signature_timestamp(
            &signed_document()?,
            &timestamp_authority(),
            &XadesTimestampConfig::new(),
        )?;
        add_xades_validation_data(
            &timestamped,
            &StaticValidationDataProvider::new()
                .with_certificate(b"chain-cert")
                .with_ocsp(b"ocsp-response"),
            &XadesValidationDataConfig::new(),
        )
    }

    #[test]
    fn xades_archive_adds_archive_timestamp_without_breaking_signature() -> XmlResult<()> {
        let archived = add_xades_archive_timestamp(
            &preserved_document()?,
            &timestamp_authority(),
            &XadesArchiveConfig::new(),
        )?;
        let report = verify_xades_archive_timestamps(
            &archived,
            &timestamp_authority(),
            &XadesArchiveConfig::new(),
        )?;
        let xades_report = verify_xades_bes_enveloped(
            &archived,
            &signing_provider(),
            &XadesConfig::new().with_signing_time("2026-06-11T12:00:00Z"),
        )?;
        let xml = to_string_compact(&archived)?;

        assert!(report.preservation_ready);
        assert_eq!(report.archive_timestamp_count, 1);
        assert_eq!(report.token_valid_count, 1);
        assert!(report.latest_token_valid);
        assert!(xades_report.valid);
        assert!(xml.contains("<xades:ArchiveTimeStamp>"));
        Ok(())
    }

    #[test]
    fn xades_archive_allows_retimestamping() -> XmlResult<()> {
        let archived = add_xades_archive_timestamp(
            &preserved_document()?,
            &timestamp_authority(),
            &XadesArchiveConfig::new(),
        )?;
        let renewed = add_xades_archive_timestamp(
            &archived,
            &timestamp_authority(),
            &XadesArchiveConfig::new(),
        )?;
        let report = verify_xades_archive_timestamps(
            &renewed,
            &timestamp_authority(),
            &XadesArchiveConfig::new(),
        )?;

        assert!(report.preservation_ready);
        assert_eq!(report.archive_timestamp_count, 2);
        assert_eq!(report.token_valid_count, 2);
        Ok(())
    }

    #[test]
    fn xades_archive_reports_missing_preservation_material() -> XmlResult<()> {
        let report = verify_xades_archive_timestamps(
            &signed_document()?,
            &timestamp_authority(),
            &XadesArchiveConfig::new(),
        )?;

        assert!(!report.signature_timestamp_present);
        assert!(!report.validation_data_present);
        assert_eq!(report.archive_timestamp_count, 0);
        assert!(!report.preservation_ready);
        Ok(())
    }

    #[test]
    fn xades_archive_rejects_add_without_validation_data() -> XmlResult<()> {
        let timestamped = add_signature_timestamp(
            &signed_document()?,
            &timestamp_authority(),
            &XadesTimestampConfig::new(),
        )?;
        let error = add_xades_archive_timestamp(
            &timestamped,
            &timestamp_authority(),
            &XadesArchiveConfig::new(),
        )
        .expect_err("archive timestamp requires validation data");

        assert_eq!(error.kind(), &ErrorKind::Signature);
        Ok(())
    }
}