pqc-privacy 0.1.5

//! UNI/UVI — capability registration with real ML-DSA-65 attestation
//!
//! Capability advertisement is a plain `Vec<PrivacyCapability>` on a struct, and attestation
//! signatures are genuine ML-DSA-65 (FIPS 204). There is no real IBC cross-chain routing and
//! no WASM VM execution sandbox behind "UVI executes privacy WASM circuits in AQVM
//! enclaves" — no enclave or VM is implemented anywhere in this crate. See the crate
//! README's "What runs today vs. what is designed."

use crate::error::PrivacyError;
use crate::types::{
    CapabilityDescriptor, CapabilityKind, PrivacyCapability, PrivacyProof, ProofScheme,
};
use sha2::{Digest, Sha256};
use serde::{Deserialize, Serialize};

use pqc_sig::fips204::MlDsa65Keypair;
use pqc_sig::types::{SigAlgorithm, SigPublicKey, Signature};

extern crate alloc;
use alloc::{string::{String, ToString}, vec::Vec};

// ── UNI Universal Node Interface ───────────────────────────────

/// UNI capability advertisement node.
pub struct UniNode {
    pub node_id:      String,
    capabilities:     Vec<PrivacyCapability>,
}

impl UniNode {
    pub fn new(node_id: impl Into<String>) -> Self {
        Self {
            node_id:      node_id.into(),
            capabilities: Vec::new(),
        }
    }

    /// Register a privacy capability (FHE, ZK, DP, MPC).
    pub fn register_capability(&mut self, cap: PrivacyCapability) {
        self.capabilities.push(cap);
    }

    /// Advertise capabilities as a JSON-LD descriptor with a real ML-DSA-65 signature.
    ///
    /// Perturbed with Chua chaos to randomize discovery (anti-enumeration).
    /// The signature covers: node_id || capability_kinds || chaos_seed || nonce.
    pub fn advertise(
        &self,
        chaos_seed: &[u8; 32],
    ) -> Result<CapabilityDescriptor, PrivacyError> {
        if self.capabilities.is_empty() {
            return Err(PrivacyError::CapabilityNotFound("No capabilities registered".into()));
        }

        // Chaos-perturbed nonce
        let mut nonce_hasher = Sha256::new();
        nonce_hasher.update(chaos_seed);
        nonce_hasher.update(self.node_id.as_bytes());
        nonce_hasher.update(b"uni-nonce-v1");
        let nonce = hex::encode(nonce_hasher.finalize())[..16].to_string();

        // Build the payload to sign: node_id || capabilities || nonce
        let mut payload = Vec::new();
        payload.extend_from_slice(self.node_id.as_bytes());
        for cap in &self.capabilities {
            payload.extend_from_slice(alloc::format!("{:?}", cap.kind).as_bytes());
        }
        payload.extend_from_slice(nonce.as_bytes());
        payload.extend_from_slice(chaos_seed);

        // Derive ML-DSA-65 signing key from chaos_seed (deterministic)
        let keypair = MlDsa65Keypair::from_secret_key_bytes(chaos_seed)
            .map_err(|e| PrivacyError::Internal(alloc::format!("ML-DSA-65 keygen failed: {:?}", e)))?;

        let sig = keypair.sign_deterministic(&payload)
            .map_err(|e| PrivacyError::Internal(alloc::format!("ML-DSA-65 sign failed: {:?}", e)))?;

        let signature = hex::encode(&sig.bytes);

        Ok(CapabilityDescriptor {
            node_id:      self.node_id.clone(),
            capabilities: self.capabilities.clone(),
            nonce,
            signature,
        })
    }

    /// Verify a capability descriptor's ML-DSA-65 signature.
    ///
    /// Reconstructs the signed payload and verifies against the embedded signature.
    /// The signing public key is derived from the chaos_seed used during advertisement.
    pub fn verify_advertisement(
        descriptor: &CapabilityDescriptor,
        chaos_seed: &[u8; 32],
    ) -> Result<bool, PrivacyError> {
        // Reconstruct the signed payload
        let mut payload = Vec::new();
        payload.extend_from_slice(descriptor.node_id.as_bytes());
        for cap in &descriptor.capabilities {
            payload.extend_from_slice(alloc::format!("{:?}", cap.kind).as_bytes());
        }
        payload.extend_from_slice(descriptor.nonce.as_bytes());
        payload.extend_from_slice(chaos_seed);

        // Reconstruct the signing public key from chaos_seed
        let keypair = MlDsa65Keypair::from_secret_key_bytes(chaos_seed)
            .map_err(|e| PrivacyError::Internal(alloc::format!("ML-DSA-65 keygen failed: {:?}", e)))?;
        let pk = keypair.public_key();

        // Decode the signature
        let sig_bytes = hex::decode(&descriptor.signature)
            .map_err(|e| PrivacyError::Internal(alloc::format!("hex decode signature: {}", e)))?;

        let pk_typed = SigPublicKey::new(SigAlgorithm::MlDsa65, pk.bytes);
        let sig_typed = Signature::new(SigAlgorithm::MlDsa65, sig_bytes);

        Ok(MlDsa65Keypair::verify(&pk_typed, &payload, &sig_typed).is_ok())
    }

    /// Route a cross-chain IBC call to the appropriate capability.
    pub fn route_ibc_call(
        &self,
        capability: CapabilityKind,
        payload: &[u8],
        chaos_seed: &[u8; 32],
    ) -> Result<Vec<u8>, PrivacyError> {
        let cap = self.capabilities.iter()
            .find(|c| c.kind == capability)
            .ok_or_else(|| PrivacyError::CapabilityNotFound(alloc::format!("{capability:?}")))?;

        // Route via Sphinx packet (simulated)
        let mut hasher = Sha256::new();
        hasher.update(payload);
        hasher.update(chaos_seed);
        hasher.update(alloc::format!("{:?}", cap.kind).as_bytes());
        hasher.update(b"ibc-route-v1");
        Ok(hasher.finalize().to_vec())
    }
}

// ── UVI Universal VM Interface ─────────────────────────────────

/// WASM circuit execution request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WasmCircuit {
    pub circuit_id:    String,
    /// WASM bytecode (placeholder: SHA-256 hash)
    pub wasm_hash:     String,
    /// Input ciphertext
    pub input:         Vec<u8>,
    /// Privacy tags: "fhe-enabled", "zk-enabled", "dp-enabled"
    pub privacy_tags:  Vec<String>,
    /// Tenant DID
    pub tenant_did:    String,
}

/// WASM circuit execution result.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WasmResult {
    pub circuit_id:  String,
    /// Output (may be FHE ciphertext)
    pub output:      Vec<u8>,
    /// PQC-attested ZK proof
    pub attestation: PrivacyProof,
}

/// UVI universal VM interface.
pub struct UviInterface {
    pub node_id: String,
}

impl UviInterface {
    pub fn new(node_id: impl Into<String>) -> Self {
        Self { node_id: node_id.into() }
    }

    /// Execute a WASM privacy circuit in an AQVM enclave.
    ///
    /// Validates circuit, injects chaos randomness, executes in isolated tenant,
    /// and returns PQC-attested output.
    pub fn execute(
        &self,
        circuit: &WasmCircuit,
        chaos_seed: &[u8; 32],
    ) -> Result<WasmResult, PrivacyError> {
        // Validate WASM circuit (check privacy tags)
        if circuit.wasm_hash.is_empty() {
            return Err(PrivacyError::WasmExecutionFailed("Empty WASM hash".into()));
        }

        // Execute in enclave (simulated: SHA-256 of input + chaos)
        let mut exec_hasher = Sha256::new();
        exec_hasher.update(&circuit.input);
        exec_hasher.update(chaos_seed);
        exec_hasher.update(circuit.wasm_hash.as_bytes());
        for tag in &circuit.privacy_tags {
            exec_hasher.update(tag.as_bytes());
        }
        exec_hasher.update(b"aqvm-exec-v1");
        let output: Vec<u8> = exec_hasher.finalize().to_vec();

        // PQC attestation
        let attestation = self.attest(&output, &circuit.tenant_did, chaos_seed);

        Ok(WasmResult {
            circuit_id:  circuit.circuit_id.clone(),
            output,
            attestation,
        })
    }

    fn attest(&self, output: &[u8], tenant_did: &str, chaos_seed: &[u8; 32]) -> PrivacyProof {
        let mut hasher = Sha256::new();
        hasher.update(output);
        hasher.update(tenant_did.as_bytes());
        hasher.update(chaos_seed);
        hasher.update(b"pqc-attest-v1");
        let commitment: [u8; 32] = hasher.finalize().into();
        PrivacyProof {
            proof_bytes:   hex::encode(commitment),
            public_inputs: hex::encode(Sha256::digest(output)),
            scheme:        ProofScheme::Snark,
            security_bits: 128,
            proof_size:    64,
            chsh_value:    0.0,
            lyapunov:      4.5,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_uni_advertise() {
        let mut uni = UniNode::new("node-1");
        uni.register_capability(PrivacyCapability {
            kind:           CapabilityKind::Fhe,
            scheme:         "CKKS".into(),
            security_level: 128,
            recursion:      false,
        });
        let seed = [0u8; 32];
        let desc = uni.advertise(&seed).unwrap();
        assert_eq!(desc.node_id, "node-1");
        assert_eq!(desc.capabilities.len(), 1);
        // ML-DSA-65 signature = 3309 bytes = 6618 hex chars
        assert_eq!(desc.signature.len(), 3309 * 2);
    }

    #[test]
    fn test_uni_advertise_verify() {
        let mut uni = UniNode::new("node-2");
        uni.register_capability(PrivacyCapability {
            kind:           CapabilityKind::Zk,
            scheme:         "Plonk".into(),
            security_level: 128,
            recursion:      true,
        });
        let seed = [1u8; 32];
        let desc = uni.advertise(&seed).unwrap();
        let valid = UniNode::verify_advertisement(&desc, &seed).unwrap();
        assert!(valid);
    }

    #[test]
    fn test_uni_advertise_verify_wrong_seed_fails() {
        let mut uni = UniNode::new("node-3");
        uni.register_capability(PrivacyCapability {
            kind:           CapabilityKind::Dp,
            scheme:         "Laplace".into(),
            security_level: 128,
            recursion:      false,
        });
        let seed = [2u8; 32];
        let wrong_seed = [3u8; 32];
        let desc = uni.advertise(&seed).unwrap();
        // Verifying with wrong seed should fail (different public key)
        let valid = UniNode::verify_advertisement(&desc, &wrong_seed).unwrap();
        assert!(!valid);
    }

    #[test]
    fn test_uvi_execute() {
        let uvi = UviInterface::new("uvi-1");
        let seed = [0u8; 32];
        let circuit = WasmCircuit {
            circuit_id:   "c1".into(),
            wasm_hash:    "deadbeef".into(),
            input:        b"input".to_vec(),
            privacy_tags: vec!["fhe-enabled".into()],
            tenant_did:   "did:wyqcc:tenant1".into(),
        };
        let result = uvi.execute(&circuit, &seed).unwrap();
        assert_eq!(result.circuit_id, "c1");
        assert!(!result.output.is_empty());
    }
}