src/pipeline/proof.rs

use chrono::{DateTime, Utc};
use hex::ToHex;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::{models::{ConsentHash, TriageRequest}, shield::redact::PiiMatch};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RedactionProof {
    pub version: String,
    pub public_commitment: String,
    pub proof: String,
    pub verified: bool,
    pub issued_at: DateTime<Utc>,
}

pub fn build_redaction_proof(request: &TriageRequest, redacted: &str, pii_map: &[PiiMatch]) -> RedactionProof {
    let payload = canonical_payload(request, redacted, pii_map);
    let public_commitment = sha256_hex(&payload);
    let proof_material = format!("{}|{}|{}", public_commitment, request.consent_hash.as_str(), redacted);
    let proof = format!("zk-redaction-sim.v1.{}", sha256_hex(proof_material.as_bytes()));
    let verified = verify_redaction_proof(request, redacted, pii_map, &public_commitment, &proof);

    RedactionProof {
        version: "zk-redaction-sim.v1".to_string(),
        public_commitment,
        proof,
        verified,
        issued_at: Utc::now(),
    }
}

pub fn verify_redaction_proof(
    request: &TriageRequest,
    redacted: &str,
    pii_map: &[PiiMatch],
    public_commitment: &str,
    proof: &str,
) -> bool {
    let expected_commitment = sha256_hex(canonical_payload(request, redacted, pii_map));
    let expected_proof = format!(
        "zk-redaction-sim.v1.{}",
        sha256_hex(format!("{}|{}|{}", expected_commitment, request.consent_hash.as_str(), redacted).as_bytes())
    );
    expected_commitment == public_commitment && expected_proof == proof
}

fn canonical_payload(request: &TriageRequest, redacted: &str, pii_map: &[PiiMatch]) -> Vec<u8> {
    let mut clone = pii_map.to_vec();
    clone.sort_by(|a, b| a.placeholder.cmp(&b.placeholder));
    let consent = ConsentHash::new(request.consent_hash.as_str().to_string())
        .map(|hash| hash.as_str().to_string())
        .unwrap_or_else(|_| request.consent_hash.as_str().to_string());
    let mut transcript = String::new();
    transcript.push_str(request.patient_id.as_str());
    transcript.push('|');
    transcript.push_str(&request.reason);
    transcript.push('|');
    transcript.push_str(&request.note);
    transcript.push('|');
    transcript.push_str(&consent);
    transcript.push('|');
    transcript.push_str(redacted);
    transcript.push('|');
    for item in clone {
        transcript.push_str(&item.entity_type);
        transcript.push('|');
        transcript.push_str(&item.original);
        transcript.push('|');
        transcript.push_str(&item.placeholder);
        transcript.push('|');
    }
    transcript.into_bytes()
}

fn sha256_hex(bytes: impl AsRef<[u8]>) -> String {
    let mut hasher = Sha256::new();
    hasher.update(bytes.as_ref());
    hasher.finalize().encode_hex::<String>()
}