mielin_cells/security/
encryption.rs

1//! State snapshot encryption
2//!
3//! This module provides encryption/decryption for agent state snapshots using AES-256-GCM.
4
5use crate::CellError;
6use ring::aead::{
7    Aad, BoundKey, Nonce, NonceSequence, OpeningKey, SealingKey, UnboundKey, AES_256_GCM,
8};
9use ring::error::Unspecified;
10use ring::rand::SecureRandom;
11use serde::{Deserialize, Serialize};
12
13/// Encryption key for state snapshots
14#[derive(Clone)]
15pub struct EncryptionKey {
16    /// Key material (32 bytes for AES-256)
17    key_bytes: Vec<u8>,
18}
19
20impl EncryptionKey {
21    /// Generate a new random encryption key
22    pub fn generate() -> Self {
23        let rng = ring::rand::SystemRandom::new();
24        let mut key_bytes = vec![0u8; 32]; // AES-256 requires 32 bytes
25        rng.fill(&mut key_bytes)
26            .expect("Failed to generate random key");
27
28        Self { key_bytes }
29    }
30
31    /// Create a key from existing bytes
32    pub fn from_bytes(key_bytes: Vec<u8>) -> Result<Self, CellError> {
33        if key_bytes.len() != 32 {
34            return Err(CellError::InvalidState(format!(
35                "Key must be 32 bytes, got {}",
36                key_bytes.len()
37            )));
38        }
39        Ok(Self { key_bytes })
40    }
41
42    /// Get the key bytes
43    pub fn as_bytes(&self) -> &[u8] {
44        &self.key_bytes
45    }
46
47    /// Export key as hex string
48    pub fn to_hex(&self) -> String {
49        self.key_bytes
50            .iter()
51            .map(|b| format!("{:02x}", b))
52            .collect()
53    }
54
55    /// Import key from hex string
56    pub fn from_hex(hex: &str) -> Result<Self, CellError> {
57        if hex.len() != 64 {
58            return Err(CellError::InvalidState(
59                "Hex string must be 64 characters (32 bytes)".to_string(),
60            ));
61        }
62
63        let mut key_bytes = Vec::with_capacity(32);
64        for i in (0..hex.len()).step_by(2) {
65            let byte_str = &hex[i..i + 2];
66            let byte = u8::from_str_radix(byte_str, 16)
67                .map_err(|e| CellError::InvalidState(format!("Invalid hex: {}", e)))?;
68            key_bytes.push(byte);
69        }
70
71        Ok(Self { key_bytes })
72    }
73}
74
75impl std::fmt::Debug for EncryptionKey {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        f.debug_struct("EncryptionKey")
78            .field("key_bytes", &"[REDACTED]")
79            .finish()
80    }
81}
82
83/// Encrypted snapshot containing encrypted state data
84#[derive(Debug, Clone, Serialize, Deserialize)]
85pub struct EncryptedSnapshot {
86    /// Agent ID
87    pub agent_id: [u8; 16],
88    /// Encrypted data
89    pub ciphertext: Vec<u8>,
90    /// Nonce used for encryption (must be unique per encryption)
91    pub nonce: Vec<u8>,
92    /// Authentication tag (included in ciphertext with AES-GCM)
93    pub timestamp: u64,
94    /// Metadata about encryption
95    pub metadata: EncryptionMetadata,
96}
97
98/// Metadata about the encryption
99#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct EncryptionMetadata {
101    /// Algorithm used
102    pub algorithm: String,
103    /// Original size before encryption
104    pub original_size: usize,
105    /// Encrypted at timestamp
106    pub encrypted_at: u64,
107}
108
109/// State encryptor for encrypting/decrypting agent state
110pub struct StateEncryptor {
111    /// Encryption key
112    key: EncryptionKey,
113}
114
115impl StateEncryptor {
116    /// Create a new state encryptor with a generated key
117    pub fn new() -> Self {
118        Self {
119            key: EncryptionKey::generate(),
120        }
121    }
122
123    /// Create a state encryptor with a specific key
124    pub fn with_key(key: EncryptionKey) -> Self {
125        Self { key }
126    }
127
128    /// Get the encryption key
129    pub fn key(&self) -> &EncryptionKey {
130        &self.key
131    }
132
133    /// Encrypt a state snapshot
134    pub fn encrypt(
135        &self,
136        agent_id: [u8; 16],
137        plaintext: &[u8],
138    ) -> Result<EncryptedSnapshot, CellError> {
139        // Generate a unique nonce for this encryption
140        let rng = ring::rand::SystemRandom::new();
141        let mut nonce_bytes = [0u8; 12];
142        rng.fill(&mut nonce_bytes)
143            .map_err(|_| CellError::InvalidState("Failed to generate nonce".to_string()))?;
144
145        // Create the sealing key
146        let unbound_key = UnboundKey::new(&AES_256_GCM, self.key.as_bytes())
147            .map_err(|_| CellError::InvalidState("Failed to create encryption key".to_string()))?;
148
149        struct FixedNonce([u8; 12]);
150        impl NonceSequence for FixedNonce {
151            fn advance(&mut self) -> Result<Nonce, Unspecified> {
152                Nonce::try_assume_unique_for_key(&self.0)
153            }
154        }
155
156        let mut sealing_key = SealingKey::new(unbound_key, FixedNonce(nonce_bytes));
157
158        // Prepare data for encryption
159        let mut in_out = plaintext.to_vec();
160
161        // Encrypt the data
162        let aad = Aad::from(&agent_id);
163        sealing_key
164            .seal_in_place_append_tag(aad, &mut in_out)
165            .map_err(|_| CellError::InvalidState("Encryption failed".to_string()))?;
166
167        let timestamp = std::time::SystemTime::now()
168            .duration_since(std::time::UNIX_EPOCH)
169            .map_err(|_| CellError::InvalidState("System time error".to_string()))?
170            .as_secs();
171
172        Ok(EncryptedSnapshot {
173            agent_id,
174            ciphertext: in_out,
175            nonce: nonce_bytes.to_vec(),
176            timestamp,
177            metadata: EncryptionMetadata {
178                algorithm: "AES-256-GCM".to_string(),
179                original_size: plaintext.len(),
180                encrypted_at: timestamp,
181            },
182        })
183    }
184
185    /// Decrypt an encrypted snapshot
186    pub fn decrypt(&self, snapshot: &EncryptedSnapshot) -> Result<Vec<u8>, CellError> {
187        // Verify nonce length
188        if snapshot.nonce.len() != 12 {
189            return Err(CellError::InvalidState("Invalid nonce length".to_string()));
190        }
191
192        let mut nonce_bytes = [0u8; 12];
193        nonce_bytes.copy_from_slice(&snapshot.nonce);
194
195        // Create the opening key
196        let unbound_key = UnboundKey::new(&AES_256_GCM, self.key.as_bytes())
197            .map_err(|_| CellError::InvalidState("Failed to create decryption key".to_string()))?;
198
199        struct FixedNonce([u8; 12]);
200        impl NonceSequence for FixedNonce {
201            fn advance(&mut self) -> Result<Nonce, Unspecified> {
202                Nonce::try_assume_unique_for_key(&self.0)
203            }
204        }
205
206        let mut opening_key = OpeningKey::new(unbound_key, FixedNonce(nonce_bytes));
207
208        // Prepare data for decryption
209        let mut in_out = snapshot.ciphertext.clone();
210
211        // Decrypt the data
212        let aad = Aad::from(&snapshot.agent_id);
213        let plaintext = opening_key
214            .open_in_place(aad, &mut in_out)
215            .map_err(|_| CellError::InvalidState("Decryption failed".to_string()))?;
216
217        Ok(plaintext.to_vec())
218    }
219
220    /// Rotate the encryption key
221    pub fn rotate_key(&mut self) -> EncryptionKey {
222        let old_key = self.key.clone();
223        self.key = EncryptionKey::generate();
224        old_key
225    }
226}
227
228impl Default for StateEncryptor {
229    fn default() -> Self {
230        Self::new()
231    }
232}
233
234#[cfg(test)]
235mod tests {
236    use super::*;
237
238    #[test]
239    fn test_encryption_key_generation() {
240        let key = EncryptionKey::generate();
241        assert_eq!(key.as_bytes().len(), 32);
242    }
243
244    #[test]
245    fn test_encryption_key_from_bytes() {
246        let key_bytes = vec![0u8; 32];
247        let key = EncryptionKey::from_bytes(key_bytes).expect("Failed to create key");
248        assert_eq!(key.as_bytes().len(), 32);
249    }
250
251    #[test]
252    fn test_encryption_key_from_bytes_invalid() {
253        let key_bytes = vec![0u8; 16]; // Wrong size
254        let result = EncryptionKey::from_bytes(key_bytes);
255        assert!(result.is_err());
256    }
257
258    #[test]
259    fn test_encryption_key_hex() {
260        let key = EncryptionKey::generate();
261        let hex = key.to_hex();
262        assert_eq!(hex.len(), 64);
263
264        let key2 = EncryptionKey::from_hex(&hex).expect("Failed to parse hex");
265        assert_eq!(key.as_bytes(), key2.as_bytes());
266    }
267
268    #[test]
269    fn test_state_encryptor_encrypt_decrypt() {
270        let encryptor = StateEncryptor::new();
271        let agent_id = [1u8; 16];
272        let plaintext = b"Hello, MielinOS! This is a secret state.";
273
274        let encrypted = encryptor
275            .encrypt(agent_id, plaintext)
276            .expect("Encryption failed");
277
278        assert_eq!(encrypted.agent_id, agent_id);
279        assert_ne!(encrypted.ciphertext, plaintext);
280
281        let decrypted = encryptor.decrypt(&encrypted).expect("Decryption failed");
282
283        assert_eq!(decrypted, plaintext);
284    }
285
286    #[test]
287    fn test_state_encryptor_decrypt_wrong_key() {
288        let encryptor1 = StateEncryptor::new();
289        let encryptor2 = StateEncryptor::new();
290
291        let agent_id = [1u8; 16];
292        let plaintext = b"Secret data";
293
294        let encrypted = encryptor1
295            .encrypt(agent_id, plaintext)
296            .expect("Encryption failed");
297
298        // Try to decrypt with wrong key
299        let result = encryptor2.decrypt(&encrypted);
300        assert!(result.is_err());
301    }
302
303    #[test]
304    fn test_state_encryptor_with_key() {
305        let key = EncryptionKey::generate();
306        let encryptor1 = StateEncryptor::with_key(key.clone());
307        let encryptor2 = StateEncryptor::with_key(key);
308
309        let agent_id = [1u8; 16];
310        let plaintext = b"Shared key test";
311
312        let encrypted = encryptor1
313            .encrypt(agent_id, plaintext)
314            .expect("Encryption failed");
315
316        let decrypted = encryptor2.decrypt(&encrypted).expect("Decryption failed");
317
318        assert_eq!(decrypted, plaintext);
319    }
320
321    #[test]
322    fn test_state_encryptor_large_data() {
323        let encryptor = StateEncryptor::new();
324        let agent_id = [1u8; 16];
325        let plaintext = vec![42u8; 1_000_000]; // 1 MB of data
326
327        let encrypted = encryptor
328            .encrypt(agent_id, &plaintext)
329            .expect("Encryption failed");
330
331        assert_eq!(encrypted.metadata.original_size, 1_000_000);
332
333        let decrypted = encryptor.decrypt(&encrypted).expect("Decryption failed");
334
335        assert_eq!(decrypted, plaintext);
336    }
337
338    #[test]
339    fn test_state_encryptor_rotate_key() {
340        let mut encryptor = StateEncryptor::new();
341        let old_key_bytes = encryptor.key().as_bytes().to_vec();
342
343        let old_key = encryptor.rotate_key();
344
345        assert_eq!(old_key.as_bytes(), old_key_bytes.as_slice());
346        assert_ne!(encryptor.key().as_bytes(), old_key_bytes.as_slice());
347    }
348
349    #[test]
350    fn test_encrypted_snapshot_metadata() {
351        let encryptor = StateEncryptor::new();
352        let agent_id = [1u8; 16];
353        let plaintext = b"Test data";
354
355        let encrypted = encryptor
356            .encrypt(agent_id, plaintext)
357            .expect("Encryption failed");
358
359        assert_eq!(encrypted.metadata.algorithm, "AES-256-GCM");
360        assert_eq!(encrypted.metadata.original_size, plaintext.len());
361        assert!(encrypted.metadata.encrypted_at > 0);
362    }
363}
mielin_cells/security/encryption.rs

mielin_cells/security/
encryption.rs
