pqc-privacy 0.1.5

//! Messenger — real AES-GCM-256 message encryption, no P2P transport
//!
//! Message content is genuinely encrypted with AES-GCM-256, keyed via real HKDF-SHA256.
//! "P2P direct mode" is a naming choice, not a transport: there is no peer-to-peer
//! networking anywhere in this crate (see [`crate::mesh`]), so "metadata obliteration via
//! DW3B" describes a design target for a mesh that does not exist yet, not a property of
//! this module. See the crate README's "What runs today vs. what is designed."

use crate::error::PrivacyError;
use crate::keyhop::hkdf_derive;
use crate::types::{PrivacyProof, ProofScheme};
use sha2::{Digest, Sha256};
use serde::{Deserialize, Serialize};
use aes_gcm::{Aes256Gcm, Key, Nonce, AeadInPlace, KeyInit};

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

/// AES-GCM-256 nonce length in bytes.
const NONCE_LEN: usize = 12;
/// AES-GCM-256 tag length in bytes.
const TAG_LEN: usize = 16;

/// Message mode selector.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MessageMode {
    /// P2P direct (1:1) with QFKH
    P2pDirect,
    /// Hybrid relay (groups > 2)
    HybridRelay,
}

/// An encrypted sovereign message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SovereignMessage {
    pub sender_did:    String,
    pub recipient_did: String,
    /// AES-GCM-256 encrypted content: nonce (12) || ciphertext || tag (16)
    pub content:       Vec<u8>,
    /// QFKH session key ID
    pub key_id:        String,
    /// ZK integrity proof
    pub proof:         PrivacyProof,
    pub mode:          MessageMode,
    pub timestamp_ms:  u64,
    /// Metadata obliterated flag
    pub metadata_free: bool,
}

/// Sovereign messenger engine.
pub struct SovereignMessenger {
    /// Participant count threshold for hybrid mode
    group_threshold: usize,
}

impl SovereignMessenger {
    pub fn new() -> Self {
        Self { group_threshold: 2 }
    }

    /// Send a message with automatic mode selection.
    ///
    /// 1:1 → P2P direct with QFKH (AES-GCM-256, key from HKDF(shared_secret, sender||recipient, "dm-enc-v1"))
    /// Group → Hybrid relay (AES-GCM-256, key from HKDF(group_key, group_id||timestamp, "group-enc-v1"))
    pub fn send(
        &self,
        sender_did: impl Into<String>,
        recipient_did: impl Into<String>,
        plaintext: &[u8],
        chaos_seed: &[u8; 32],
        timestamp_ms: u64,
        participant_count: usize,
    ) -> Result<SovereignMessage, PrivacyError> {
        let sender_did    = sender_did.into();
        let recipient_did = recipient_did.into();

        let mode = if participant_count <= self.group_threshold {
            MessageMode::P2pDirect
        } else {
            MessageMode::HybridRelay
        };

        // Encrypt content with AES-GCM-256
        let content = match mode {
            MessageMode::P2pDirect => {
                self.encrypt_dm(plaintext, chaos_seed, &sender_did, &recipient_did)?
            }
            MessageMode::HybridRelay => {
                self.encrypt_group(plaintext, chaos_seed, &recipient_did, timestamp_ms)?
            }
        };

        // QFKH key ID
        let mut key_hasher = Sha256::new();
        key_hasher.update(chaos_seed);
        key_hasher.update(timestamp_ms.to_le_bytes());
        key_hasher.update(b"qfkh-key-id");
        let key_id = hex::encode(key_hasher.finalize())[..16].to_string();

        // ZK integrity proof
        let proof = self.prove_integrity(&content, chaos_seed);

        Ok(SovereignMessage {
            sender_did,
            recipient_did,
            content,
            key_id,
            proof,
            mode,
            timestamp_ms,
            metadata_free: true,
        })
    }

    /// Decrypt a received direct message.
    pub fn receive(
        &self,
        msg: &SovereignMessage,
        chaos_seed: &[u8; 32],
    ) -> Result<Vec<u8>, PrivacyError> {
        match msg.mode {
            MessageMode::P2pDirect => {
                self.decrypt_dm(&msg.content, chaos_seed, &msg.sender_did, &msg.recipient_did)
            }
            MessageMode::HybridRelay => {
                self.decrypt_group(&msg.content, chaos_seed, &msg.recipient_did, msg.timestamp_ms)
            }
        }
    }

    /// Obliterate metadata from a message.
    ///
    /// Strips IP, timestamps, and replaces with chaos-randomized values.
    pub fn obliterate_metadata(
        &self,
        msg: &mut SovereignMessage,
        chaos_seed: &[u8; 32],
    ) {
        // Replace timestamp with chaos-randomized value
        let seed_val = u64::from_le_bytes(chaos_seed[..8].try_into().unwrap_or([0u8; 8]));
        msg.timestamp_ms = seed_val % 1_000_000_000_000;
        msg.metadata_free = true;
    }

    // ── Private helpers ───────────────────────────────────────────────────────

    /// Encrypt a direct message with AES-GCM-256.
    ///
    /// Key = HKDF(chaos_seed, sender_id || recipient_id, "dm-enc-v1")
    /// Nonce = first 12 bytes of SHA-256(sender_id || recipient_id || "dm-nonce-v1")
    fn encrypt_dm(
        &self,
        plaintext: &[u8],
        chaos_seed: &[u8; 32],
        sender_id: &str,
        recipient_id: &str,
    ) -> Result<Vec<u8>, PrivacyError> {
        let key = self.derive_dm_key(chaos_seed, sender_id, recipient_id)?;
        let nonce = self.derive_dm_nonce(sender_id, recipient_id);
        self.aes_gcm_encrypt(plaintext, &key, &nonce, sender_id.as_bytes())
    }

    /// Decrypt a direct message with AES-GCM-256.
    fn decrypt_dm(
        &self,
        ciphertext: &[u8],
        chaos_seed: &[u8; 32],
        sender_id: &str,
        recipient_id: &str,
    ) -> Result<Vec<u8>, PrivacyError> {
        let key = self.derive_dm_key(chaos_seed, sender_id, recipient_id)?;
        let nonce = self.derive_dm_nonce(sender_id, recipient_id);
        self.aes_gcm_decrypt(ciphertext, &key, &nonce, sender_id.as_bytes())
    }

    /// Encrypt a group message with AES-GCM-256.
    ///
    /// Key = HKDF(chaos_seed, group_id || timestamp, "group-enc-v1")
    /// Nonce = first 12 bytes of SHA-256(group_id || timestamp || "group-nonce-v1")
    fn encrypt_group(
        &self,
        plaintext: &[u8],
        chaos_seed: &[u8; 32],
        group_id: &str,
        timestamp_ms: u64,
    ) -> Result<Vec<u8>, PrivacyError> {
        let key = self.derive_group_key(chaos_seed, group_id, timestamp_ms)?;
        let nonce = self.derive_group_nonce(group_id, timestamp_ms);
        self.aes_gcm_encrypt(plaintext, &key, &nonce, group_id.as_bytes())
    }

    /// Decrypt a group message with AES-GCM-256.
    fn decrypt_group(
        &self,
        ciphertext: &[u8],
        chaos_seed: &[u8; 32],
        group_id: &str,
        timestamp_ms: u64,
    ) -> Result<Vec<u8>, PrivacyError> {
        let key = self.derive_group_key(chaos_seed, group_id, timestamp_ms)?;
        let nonce = self.derive_group_nonce(group_id, timestamp_ms);
        self.aes_gcm_decrypt(ciphertext, &key, &nonce, group_id.as_bytes())
    }

    fn derive_dm_key(
        &self,
        chaos_seed: &[u8; 32],
        sender_id: &str,
        recipient_id: &str,
    ) -> Result<[u8; 32], PrivacyError> {
        let mut salt = Vec::with_capacity(sender_id.len() + recipient_id.len());
        salt.extend_from_slice(sender_id.as_bytes());
        salt.extend_from_slice(recipient_id.as_bytes());
        let mut key = [0u8; 32];
        hkdf_derive(chaos_seed, &salt, b"dm-enc-v1", &mut key)
            .map_err(|e| PrivacyError::P2pConnectionFailed(alloc::format!("HKDF DM key: {:?}", e)))?;
        Ok(key)
    }

    fn derive_dm_nonce(&self, sender_id: &str, recipient_id: &str) -> [u8; NONCE_LEN] {
        let mut hasher = Sha256::new();
        hasher.update(sender_id.as_bytes());
        hasher.update(recipient_id.as_bytes());
        hasher.update(b"dm-nonce-v1");
        let h: [u8; 32] = hasher.finalize().into();
        h[..NONCE_LEN].try_into().unwrap_or([0u8; NONCE_LEN])
    }

    fn derive_group_key(
        &self,
        chaos_seed: &[u8; 32],
        group_id: &str,
        timestamp_ms: u64,
    ) -> Result<[u8; 32], PrivacyError> {
        let mut salt = Vec::with_capacity(group_id.len() + 8);
        salt.extend_from_slice(group_id.as_bytes());
        salt.extend_from_slice(&timestamp_ms.to_le_bytes());
        let mut key = [0u8; 32];
        hkdf_derive(chaos_seed, &salt, b"group-enc-v1", &mut key)
            .map_err(|e| PrivacyError::P2pConnectionFailed(alloc::format!("HKDF group key: {:?}", e)))?;
        Ok(key)
    }

    fn derive_group_nonce(&self, group_id: &str, timestamp_ms: u64) -> [u8; NONCE_LEN] {
        let mut hasher = Sha256::new();
        hasher.update(group_id.as_bytes());
        hasher.update(timestamp_ms.to_le_bytes());
        hasher.update(b"group-nonce-v1");
        let h: [u8; 32] = hasher.finalize().into();
        h[..NONCE_LEN].try_into().unwrap_or([0u8; NONCE_LEN])
    }

    fn aes_gcm_encrypt(
        &self,
        plaintext: &[u8],
        key: &[u8; 32],
        nonce_bytes: &[u8; NONCE_LEN],
        aad: &[u8],
    ) -> Result<Vec<u8>, PrivacyError> {
        let cipher_key = Key::<Aes256Gcm>::from_slice(key);
        let cipher = Aes256Gcm::new(cipher_key);
        let nonce = Nonce::from_slice(nonce_bytes);

        let mut buffer = plaintext.to_vec();
        buffer.reserve(TAG_LEN);
        cipher.encrypt_in_place(nonce, aad, &mut buffer)
            .map_err(|e| PrivacyError::P2pConnectionFailed(alloc::format!("AES-GCM encrypt: {:?}", e)))?;

        // Prepend nonce
        let mut output = Vec::with_capacity(NONCE_LEN + buffer.len());
        output.extend_from_slice(nonce_bytes);
        output.extend_from_slice(&buffer);
        Ok(output)
    }

    fn aes_gcm_decrypt(
        &self,
        ciphertext: &[u8],
        key: &[u8; 32],
        nonce_bytes: &[u8; NONCE_LEN],
        aad: &[u8],
    ) -> Result<Vec<u8>, PrivacyError> {
        if ciphertext.len() < NONCE_LEN + TAG_LEN {
            return Err(PrivacyError::P2pConnectionFailed("Ciphertext too short".into()));
        }

        // Skip the prepended nonce (we use the derived nonce for decryption)
        let ct = &ciphertext[NONCE_LEN..];

        let cipher_key = Key::<Aes256Gcm>::from_slice(key);
        let cipher = Aes256Gcm::new(cipher_key);
        let nonce = Nonce::from_slice(nonce_bytes);

        let mut buffer = ct.to_vec();
        cipher.decrypt_in_place(nonce, aad, &mut buffer)
            .map_err(|e| PrivacyError::P2pConnectionFailed(alloc::format!("AES-GCM decrypt: {:?}", e)))?;

        Ok(buffer)
    }

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

impl Default for SovereignMessenger {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_send_receive_dm() {
        let messenger = SovereignMessenger::new();
        let seed = [0u8; 32];
        let msg = messenger.send("did:alice", "did:bob", b"hello", &seed, 1000, 1).unwrap();
        assert_eq!(msg.mode, MessageMode::P2pDirect);
        assert!(msg.metadata_free);
        let decrypted = messenger.receive(&msg, &seed).unwrap();
        assert_eq!(decrypted, b"hello");
    }

    #[test]
    fn test_group_mode() {
        let messenger = SovereignMessenger::new();
        let seed = [0u8; 32];
        let msg = messenger.send("did:alice", "did:group", b"hi", &seed, 1000, 5).unwrap();
        assert_eq!(msg.mode, MessageMode::HybridRelay);
        let decrypted = messenger.receive(&msg, &seed).unwrap();
        assert_eq!(decrypted, b"hi");
    }

    #[test]
    fn test_wrong_seed_fails() {
        let messenger = SovereignMessenger::new();
        let seed1 = [1u8; 32];
        let seed2 = [2u8; 32];
        let msg = messenger.send("did:alice", "did:bob", b"secret", &seed1, 1000, 1).unwrap();
        // Decrypting with wrong seed should fail (AES-GCM tag mismatch)
        assert!(messenger.receive(&msg, &seed2).is_err());
    }
}