1pub use errors::*;
16pub use key_management::{FileKeyStorage, KeyStorage, KeyStore as KeyManagementStore};
17
18use crate::workspace_persistence::WorkspacePersistence;
19use aes_gcm::{
20 aead::{Aead, KeyInit},
21 Aes256Gcm, Nonce,
22};
23use argon2::{
24 password_hash::{PasswordHasher, SaltString},
25 Argon2, Params,
26};
27use base64::{engine::general_purpose, Engine as _};
28use chacha20poly1305::{ChaCha20Poly1305, Key as ChaChaKey};
29use pbkdf2::pbkdf2_hmac;
30use rand::{thread_rng, Rng};
31use serde::{Deserialize, Serialize};
32use sha2::Sha256;
33use std::fmt;
34use thiserror::Error;
35use tracing;
36
37#[cfg(target_os = "windows")]
38use windows::Win32::Security::Credentials::{
39 CredDeleteW, CredReadW, CredWriteW, CREDENTIALW, CRED_PERSIST_LOCAL_MACHINE, CRED_TYPE_GENERIC,
40};
41
42#[derive(Error, Debug)]
44pub enum EncryptionError {
45 #[error("Encryption failure: {0}")]
47 Encryption(String),
48 #[error("Decryption failure: {0}")]
50 Decryption(String),
51 #[error("Invalid key: {0}")]
53 InvalidKey(String),
54 #[error("Invalid ciphertext: {0}")]
56 InvalidCiphertext(String),
57 #[error("Key derivation failure: {0}")]
59 KeyDerivation(String),
60 #[error("Generic encryption error: {message}")]
62 Generic {
63 message: String,
65 },
66}
67
68pub type Result<T> = std::result::Result<T, EncryptionError>;
70
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
73pub enum EncryptionAlgorithm {
74 Aes256Gcm,
76 ChaCha20Poly1305,
78}
79
80impl fmt::Display for EncryptionAlgorithm {
81 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
82 match self {
83 EncryptionAlgorithm::Aes256Gcm => write!(f, "aes256-gcm"),
84 EncryptionAlgorithm::ChaCha20Poly1305 => write!(f, "chacha20-poly1305"),
85 }
86 }
87}
88
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
91pub enum KeyDerivationMethod {
92 Pbkdf2,
94 Argon2,
96}
97
98impl fmt::Display for KeyDerivationMethod {
99 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
100 match self {
101 KeyDerivationMethod::Pbkdf2 => write!(f, "pbkdf2"),
102 KeyDerivationMethod::Argon2 => write!(f, "argon2"),
103 }
104 }
105}
106
107pub struct EncryptionKey {
109 algorithm: EncryptionAlgorithm,
110 key_data: Vec<u8>,
111}
112
113impl EncryptionKey {
114 pub fn new(algorithm: EncryptionAlgorithm, key_data: Vec<u8>) -> Result<Self> {
116 let expected_len = match algorithm {
117 EncryptionAlgorithm::Aes256Gcm => 32, EncryptionAlgorithm::ChaCha20Poly1305 => 32, };
120
121 if key_data.len() != expected_len {
122 return Err(EncryptionError::InvalidKey(format!(
123 "Key must be {} bytes for {}, got {}",
124 expected_len,
125 algorithm,
126 key_data.len()
127 )));
128 }
129
130 Ok(Self {
131 algorithm,
132 key_data,
133 })
134 }
135
136 pub fn from_password_pbkdf2(
138 password: &str,
139 salt: Option<&[u8]>,
140 algorithm: EncryptionAlgorithm,
141 ) -> Result<Self> {
142 let salt = salt
143 .map(|s| s.to_vec())
144 .unwrap_or_else(|| thread_rng().random::<[u8; 32]>().to_vec());
145
146 let mut key = vec![0u8; 32];
147 pbkdf2_hmac::<Sha256>(password.as_bytes(), &salt, 100_000, &mut key);
148
149 Self::new(algorithm, key)
150 }
151
152 pub fn from_password_argon2(
154 password: &str,
155 salt: Option<&[u8]>,
156 algorithm: EncryptionAlgorithm,
157 ) -> Result<Self> {
158 let salt_string = if let Some(salt) = salt {
159 SaltString::encode_b64(salt)
160 .map_err(|e| EncryptionError::KeyDerivation(e.to_string()))?
161 } else {
162 let mut salt_bytes = [0u8; 32];
164 thread_rng().fill(&mut salt_bytes);
165 SaltString::encode_b64(&salt_bytes)
166 .map_err(|e| EncryptionError::KeyDerivation(e.to_string()))?
167 };
168
169 let argon2 = Argon2::new(
170 argon2::Algorithm::Argon2id,
171 argon2::Version::V0x13,
172 Params::new(65536, 3, 1, Some(32))
173 .map_err(|e| EncryptionError::KeyDerivation(e.to_string()))?,
174 );
175
176 let hash = argon2
177 .hash_password(password.as_bytes(), &salt_string)
178 .map_err(|e| EncryptionError::KeyDerivation(e.to_string()))?;
179
180 let key_bytes = hash
181 .hash
182 .ok_or_else(|| EncryptionError::KeyDerivation("Failed to derive key hash".to_string()))?
183 .as_bytes()
184 .to_vec();
185 Self::new(algorithm, key_bytes)
186 }
187
188 pub fn encrypt(&self, plaintext: &str, associated_data: Option<&[u8]>) -> Result<String> {
190 match self.algorithm {
191 EncryptionAlgorithm::Aes256Gcm => self.encrypt_aes_gcm(plaintext, associated_data),
192 EncryptionAlgorithm::ChaCha20Poly1305 => {
193 self.encrypt_chacha20(plaintext, associated_data)
194 }
195 }
196 }
197
198 pub fn decrypt(&self, ciphertext: &str, associated_data: Option<&[u8]>) -> Result<String> {
200 match self.algorithm {
201 EncryptionAlgorithm::Aes256Gcm => self.decrypt_aes_gcm(ciphertext, associated_data),
202 EncryptionAlgorithm::ChaCha20Poly1305 => {
203 self.decrypt_chacha20(ciphertext, associated_data)
204 }
205 }
206 }
207
208 fn encrypt_aes_gcm(&self, plaintext: &str, associated_data: Option<&[u8]>) -> Result<String> {
209 let key_array: [u8; 32] =
211 self.key_data.as_slice().try_into().map_err(|_| {
212 EncryptionError::InvalidKey("Key length must be 32 bytes".to_string())
213 })?;
214 let cipher = Aes256Gcm::new(&key_array.into());
215 let nonce: [u8; 12] = thread_rng().random(); let nonce = Nonce::from(nonce);
217
218 let ciphertext = cipher
219 .encrypt(&nonce, plaintext.as_bytes())
220 .map_err(|e| EncryptionError::Encryption(e.to_string()))?;
221
222 let mut result = nonce.to_vec();
223 result.extend_from_slice(&ciphertext);
224
225 if let Some(aad) = associated_data {
227 result.extend_from_slice(aad);
228 }
229
230 Ok(general_purpose::STANDARD.encode(&result))
231 }
232
233 fn decrypt_aes_gcm(&self, ciphertext: &str, associated_data: Option<&[u8]>) -> Result<String> {
234 let data = general_purpose::STANDARD
235 .decode(ciphertext)
236 .map_err(|e| EncryptionError::InvalidCiphertext(e.to_string()))?;
237
238 if data.len() < 12 {
239 return Err(EncryptionError::InvalidCiphertext("Ciphertext too short".to_string()));
240 }
241
242 let nonce_array: [u8; 12] = data[0..12].try_into().map_err(|_| {
244 EncryptionError::InvalidCiphertext("Nonce must be 12 bytes".to_string())
245 })?;
246 let nonce = Nonce::from(nonce_array);
247
248 let ciphertext_len = if let Some(aad) = &associated_data {
249 let aad_len = aad.len();
251 data.len() - 12 - aad_len
252 } else {
253 data.len() - 12
254 };
255
256 let ciphertext = &data[12..12 + ciphertext_len];
257
258 let key_array: [u8; 32] =
260 self.key_data.as_slice().try_into().map_err(|_| {
261 EncryptionError::InvalidKey("Key length must be 32 bytes".to_string())
262 })?;
263 let cipher = Aes256Gcm::new(&key_array.into());
264
265 let plaintext = cipher
266 .decrypt(&nonce, ciphertext.as_ref())
267 .map_err(|e| EncryptionError::Decryption(e.to_string()))?;
268
269 String::from_utf8(plaintext)
270 .map_err(|e| EncryptionError::Decryption(format!("Invalid UTF-8: {}", e)))
271 }
272
273 pub fn encrypt_chacha20(
282 &self,
283 plaintext: &str,
284 _associated_data: Option<&[u8]>,
285 ) -> Result<String> {
286 let key = ChaChaKey::from_slice(&self.key_data);
287 let cipher = ChaCha20Poly1305::new(key);
288 let nonce: [u8; 12] = thread_rng().random(); let nonce = chacha20poly1305::Nonce::from_slice(&nonce);
290
291 let ciphertext = cipher
292 .encrypt(nonce, plaintext.as_bytes())
293 .map_err(|e| EncryptionError::Encryption(e.to_string()))?;
294
295 let mut result = nonce.to_vec();
296 result.extend_from_slice(&ciphertext);
297
298 Ok(general_purpose::STANDARD.encode(&result))
299 }
300
301 pub fn decrypt_chacha20(
310 &self,
311 ciphertext: &str,
312 _associated_data: Option<&[u8]>,
313 ) -> Result<String> {
314 let data = general_purpose::STANDARD
315 .decode(ciphertext)
316 .map_err(|e| EncryptionError::InvalidCiphertext(e.to_string()))?;
317
318 if data.len() < 12 {
319 return Err(EncryptionError::InvalidCiphertext("Ciphertext too short".to_string()));
320 }
321
322 let nonce = chacha20poly1305::Nonce::from_slice(&data[0..12]);
323 let ciphertext_data = &data[12..];
324 let key = ChaChaKey::from_slice(&self.key_data);
325 let cipher = ChaCha20Poly1305::new(key);
326
327 let plaintext = cipher
328 .decrypt(nonce, ciphertext_data.as_ref())
329 .map_err(|e| EncryptionError::Decryption(e.to_string()))?;
330
331 String::from_utf8(plaintext)
332 .map_err(|e| EncryptionError::Decryption(format!("Invalid UTF-8: {}", e)))
333 }
334}
335
336pub struct KeyStore {
338 keys: std::collections::HashMap<String, EncryptionKey>,
339}
340
341impl KeyStore {
342 pub fn new() -> Self {
344 Self {
345 keys: std::collections::HashMap::new(),
346 }
347 }
348
349 pub fn store_key(&mut self, id: String, key: EncryptionKey) {
351 self.keys.insert(id, key);
352 }
353
354 pub fn get_key(&self, id: &str) -> Option<&EncryptionKey> {
356 self.keys.get(id)
357 }
358
359 pub fn remove_key(&mut self, id: &str) -> bool {
361 self.keys.remove(id).is_some()
362 }
363
364 pub fn list_keys(&self) -> Vec<String> {
366 self.keys.keys().cloned().collect()
367 }
368
369 pub fn derive_and_store_key(
371 &mut self,
372 id: String,
373 password: &str,
374 algorithm: EncryptionAlgorithm,
375 method: KeyDerivationMethod,
376 ) -> Result<()> {
377 let key = match method {
378 KeyDerivationMethod::Pbkdf2 => {
379 EncryptionKey::from_password_pbkdf2(password, None, algorithm)?
380 }
381 KeyDerivationMethod::Argon2 => {
382 EncryptionKey::from_password_argon2(password, None, algorithm)?
383 }
384 };
385 self.store_key(id, key);
386 Ok(())
387 }
388}
389
390impl Default for KeyStore {
391 fn default() -> Self {
392 Self::new()
393 }
394}
395
396static KEY_STORE: once_cell::sync::OnceCell<KeyStore> = once_cell::sync::OnceCell::new();
398
399pub fn init_key_store() -> &'static KeyStore {
408 KEY_STORE.get_or_init(KeyStore::default)
409}
410
411pub fn get_key_store() -> Option<&'static KeyStore> {
418 KEY_STORE.get()
419}
420
421pub struct MasterKeyManager {
428 _service_name: String,
430 _account_name: String,
432}
433
434impl MasterKeyManager {
435 pub fn new() -> Self {
437 Self {
438 _service_name: "com.mockforge.encryption".to_string(),
439 _account_name: "master_key".to_string(),
440 }
441 }
442
443 pub fn generate_master_key(&self) -> Result<()> {
445 let master_key_bytes: [u8; 32] = rand::random();
446 let master_key_b64 = general_purpose::STANDARD.encode(master_key_bytes);
447
448 #[cfg(target_os = "macos")]
451 {
452 self.store_in_macos_keychain(&master_key_b64)?;
454 }
455 #[cfg(target_os = "linux")]
456 {
457 self.store_in_linux_keyring(&master_key_b64)?;
459 }
460 #[cfg(target_os = "windows")]
461 {
462 self.store_in_windows_credential_manager(&master_key_b64)?;
464 }
465
466 Ok(())
467 }
468
469 pub fn get_master_key(&self) -> Result<EncryptionKey> {
471 let master_key_b64 = self.retrieve_from_keychain()?;
472 let master_key_bytes = general_purpose::STANDARD
473 .decode(master_key_b64)
474 .map_err(|e| EncryptionError::InvalidKey(e.to_string()))?;
475
476 if master_key_bytes.len() != 32 {
477 return Err(EncryptionError::InvalidKey("Invalid master key length".to_string()));
478 }
479
480 EncryptionKey::new(EncryptionAlgorithm::ChaCha20Poly1305, master_key_bytes)
481 }
482
483 pub fn has_master_key(&self) -> bool {
485 self.retrieve_from_keychain().is_ok()
486 }
487
488 #[cfg(target_os = "macos")]
490 fn store_in_macos_keychain(&self, key: &str) -> Result<()> {
491 use std::os::unix::fs::PermissionsExt;
492
493 let home = std::env::var("HOME").map_err(|_| {
494 EncryptionError::InvalidKey("HOME environment variable not set".to_string())
495 })?;
496 let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
497
498 if let Some(parent) = key_path.parent() {
500 std::fs::create_dir_all(parent).map_err(|e| {
501 EncryptionError::InvalidKey(format!("Failed to create directory: {}", e))
502 })?;
503 }
504
505 std::fs::write(&key_path, key).map_err(|e| {
507 EncryptionError::InvalidKey(format!("Failed to write master key: {}", e))
508 })?;
509
510 let mut perms = std::fs::metadata(&key_path)
512 .map_err(|e| EncryptionError::InvalidKey(format!("Failed to get metadata: {}", e)))?
513 .permissions();
514 perms.set_mode(0o600);
515 std::fs::set_permissions(&key_path, perms).map_err(|e| {
516 EncryptionError::InvalidKey(format!("Failed to set permissions: {}", e))
517 })?;
518
519 Ok(())
520 }
521
522 #[cfg(target_os = "linux")]
523 fn store_in_linux_keyring(&self, key: &str) -> Result<()> {
524 use std::os::unix::fs::PermissionsExt;
525
526 let home = std::env::var("HOME").map_err(|_| {
527 EncryptionError::InvalidKey("HOME environment variable not set".to_string())
528 })?;
529 let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
530
531 if let Some(parent) = key_path.parent() {
533 std::fs::create_dir_all(parent).map_err(|e| {
534 EncryptionError::InvalidKey(format!("Failed to create directory: {}", e))
535 })?;
536 }
537
538 std::fs::write(&key_path, key).map_err(|e| {
540 EncryptionError::InvalidKey(format!("Failed to write master key: {}", e))
541 })?;
542
543 let mut perms = std::fs::metadata(&key_path)
545 .map_err(|e| EncryptionError::InvalidKey(format!("Failed to get metadata: {}", e)))?
546 .permissions();
547 perms.set_mode(0o600);
548 std::fs::set_permissions(&key_path, perms).map_err(|e| {
549 EncryptionError::InvalidKey(format!("Failed to set permissions: {}", e))
550 })?;
551
552 Ok(())
553 }
554
555 #[cfg(target_os = "windows")]
556 fn store_in_windows_credential_manager(&self, key: &str) -> Result<()> {
557 use std::ffi::OsString;
558 use std::os::windows::ffi::OsStringExt;
559 use windows::core::PCWSTR;
560 use windows::Win32::Foundation::ERROR_NO_SUCH_LOGON_SESSION;
561 use windows::Win32::Security::Credentials::{
562 CredWriteW, CREDENTIALW, CRED_PERSIST_LOCAL_MACHINE, CRED_TYPE_GENERIC,
563 };
564
565 let target_name = "MockForge/MasterKey";
566 let target_name_wide: Vec<u16> =
567 OsString::from(target_name).encode_wide().chain(std::iter::once(0)).collect();
568
569 let credential_blob: Vec<u16> =
570 OsString::from(key).encode_wide().chain(std::iter::once(0)).collect();
571
572 let mut credential = CREDENTIALW {
573 Flags: 0,
574 Type: CRED_TYPE_GENERIC,
575 TargetName: PCWSTR::from_raw(target_name_wide.as_ptr()),
576 Comment: PCWSTR::null(),
577 LastWritten: windows::Win32::Foundation::FILETIME::default(),
578 CredentialBlobSize: (credential_blob.len() * 2) as u32,
579 CredentialBlob: credential_blob.as_ptr() as *mut u8,
580 Persist: CRED_PERSIST_LOCAL_MACHINE,
581 AttributeCount: 0,
582 Attributes: std::ptr::null_mut(),
583 TargetAlias: PCWSTR::null(),
584 UserName: PCWSTR::null(),
585 };
586
587 unsafe {
591 CredWriteW(&mut credential, 0).map_err(|e| {
592 EncryptionError::InvalidKey(format!("Failed to store credential: {:?}", e))
593 })?;
594 }
595
596 Ok(())
597 }
598
599 fn retrieve_from_keychain(&self) -> Result<String> {
600 #[cfg(target_os = "macos")]
602 {
603 let home = std::env::var("HOME").map_err(|_| {
604 EncryptionError::InvalidKey("HOME environment variable not set".to_string())
605 })?;
606 let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
607 let key_str = std::fs::read_to_string(&key_path).map_err(|_| {
608 EncryptionError::InvalidKey("Master key not found in keychain".to_string())
609 })?;
610 Ok(key_str.trim().to_string())
612 }
613
614 #[cfg(target_os = "linux")]
615 {
616 let home = std::env::var("HOME").map_err(|_| {
617 EncryptionError::InvalidKey("HOME environment variable not set".to_string())
618 })?;
619 let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
620 let key_str = std::fs::read_to_string(&key_path).map_err(|_| {
621 EncryptionError::InvalidKey("Master key not found in keychain".to_string())
622 })?;
623 Ok(key_str.trim().to_string())
625 }
626
627 #[cfg(target_os = "windows")]
628 {
629 self.retrieve_from_windows_credential_manager()
631 }
632
633 #[cfg(not(any(target_os = "macos", target_os = "linux", target_os = "windows")))]
634 {
635 let key_str = std::env::var("MOCKFORGE_MASTER_KEY").map_err(|_| {
637 EncryptionError::InvalidKey("Master key not found in keychain".to_string())
638 })?;
639 Ok(key_str.trim().to_string())
641 }
642 }
643
644 #[cfg(target_os = "windows")]
645 fn retrieve_from_windows_credential_manager(&self) -> Result<String> {
646 use std::ffi::OsString;
647 use std::os::windows::ffi::OsStringExt;
648 use windows::core::PCWSTR;
649 use windows::Win32::Security::Credentials::{
650 CredFree, CredReadW, CREDENTIALW, CRED_TYPE_GENERIC,
651 };
652
653 let target_name = "MockForge/MasterKey";
654 let target_name_wide: Vec<u16> =
655 OsString::from(target_name).encode_wide().chain(std::iter::once(0)).collect();
656
657 let mut credential_ptr: *mut CREDENTIALW = std::ptr::null_mut();
658
659 unsafe {
664 CredReadW(
665 PCWSTR::from_raw(target_name_wide.as_ptr()),
666 CRED_TYPE_GENERIC,
667 0,
668 &mut credential_ptr,
669 )
670 .map_err(|e| {
671 EncryptionError::InvalidKey(format!("Failed to read credential: {:?}", e))
672 })?;
673
674 if credential_ptr.is_null() {
675 return Err(EncryptionError::InvalidKey("Credential not found".to_string()));
676 }
677
678 let credential = &*credential_ptr;
680
681 let blob_slice = std::slice::from_raw_parts(
684 credential.CredentialBlob as *const u16,
685 credential.CredentialBlobSize as usize / 2, );
687
688 let credential_str = OsString::from_wide(blob_slice)
689 .to_string_lossy()
690 .trim_end_matches('\0')
691 .to_string();
692
693 CredFree(credential_ptr as *mut std::ffi::c_void);
695
696 Ok(credential_str)
697 }
698 }
699}
700
701impl Default for MasterKeyManager {
702 fn default() -> Self {
703 Self::new()
704 }
705}
706
707pub struct WorkspaceKeyManager {
712 master_key_manager: MasterKeyManager,
714 key_storage: std::cell::RefCell<FileKeyStorage>,
716}
717
718impl WorkspaceKeyManager {
719 pub fn new() -> Self {
721 Self {
722 master_key_manager: MasterKeyManager::new(),
723 key_storage: std::cell::RefCell::new(FileKeyStorage::new()),
724 }
725 }
726
727 pub fn with_storage_path<P: AsRef<std::path::Path>>(path: P) -> Self {
729 Self {
730 master_key_manager: MasterKeyManager::new(),
731 key_storage: std::cell::RefCell::new(FileKeyStorage::with_path(path)),
732 }
733 }
734
735 pub fn generate_workspace_key(&self, workspace_id: &str) -> Result<String> {
737 let workspace_key_bytes: [u8; 32] = rand::random();
739
740 let master_key = self.master_key_manager.get_master_key()?;
742
743 let workspace_key_b64 = master_key.encrypt_chacha20(
745 &general_purpose::STANDARD.encode(workspace_key_bytes),
746 Some(workspace_id.as_bytes()),
747 )?;
748
749 self.store_workspace_key(workspace_id, &workspace_key_b64)?;
751
752 Ok(workspace_key_b64)
753 }
754
755 pub fn get_workspace_key(&self, workspace_id: &str) -> Result<EncryptionKey> {
757 let encrypted_key_b64 = self.retrieve_workspace_key(workspace_id)?;
758 let master_key = self.master_key_manager.get_master_key()?;
759
760 let decrypted_key_b64 =
761 master_key.decrypt_chacha20(&encrypted_key_b64, Some(workspace_id.as_bytes()))?;
762
763 let workspace_key_bytes = general_purpose::STANDARD
764 .decode(decrypted_key_b64)
765 .map_err(|e| EncryptionError::InvalidKey(e.to_string()))?;
766
767 if workspace_key_bytes.len() != 32 {
768 return Err(EncryptionError::InvalidKey("Invalid workspace key length".to_string()));
769 }
770
771 EncryptionKey::new(EncryptionAlgorithm::ChaCha20Poly1305, workspace_key_bytes)
772 }
773
774 pub fn has_workspace_key(&self, workspace_id: &str) -> bool {
776 self.retrieve_workspace_key(workspace_id).is_ok()
777 }
778
779 pub fn generate_workspace_key_backup(&self, workspace_id: &str) -> Result<String> {
781 let encrypted_key = self.retrieve_workspace_key(workspace_id)?;
782
783 let backup_string = self.format_backup_string(&encrypted_key);
786
787 Ok(backup_string)
788 }
789
790 pub fn restore_workspace_key_from_backup(
792 &self,
793 workspace_id: &str,
794 backup_string: &str,
795 ) -> Result<()> {
796 let encrypted_key = self.parse_backup_string(backup_string)?;
797 self.store_workspace_key(workspace_id, &encrypted_key)
798 }
799
800 fn store_workspace_key(&self, workspace_id: &str, encrypted_key: &str) -> Result<()> {
802 self.key_storage
803 .borrow_mut()
804 .store_key(&workspace_id.to_string(), encrypted_key.as_bytes())
805 .map_err(|e| {
806 EncryptionError::InvalidKey(format!("Failed to store workspace key: {:?}", e))
807 })
808 }
809
810 fn retrieve_workspace_key(&self, workspace_id: &str) -> Result<String> {
811 match self.key_storage.borrow().retrieve_key(&workspace_id.to_string()) {
813 Ok(encrypted_bytes) => String::from_utf8(encrypted_bytes).map_err(|e| {
814 EncryptionError::InvalidKey(format!("Invalid UTF-8 in stored key: {}", e))
815 }),
816 Err(_) => {
817 let old_key_file = format!("workspace_{}_key.enc", workspace_id);
819 match std::fs::read_to_string(&old_key_file) {
820 Ok(encrypted_key) => {
821 if let Err(e) = self
823 .key_storage
824 .borrow_mut()
825 .store_key(&workspace_id.to_string(), encrypted_key.as_bytes())
826 {
827 tracing::warn!(
828 "Failed to migrate workspace key to new storage: {:?}",
829 e
830 );
831 } else {
832 let _ = std::fs::remove_file(&old_key_file);
834 }
835 Ok(encrypted_key)
836 }
837 Err(_) => Err(EncryptionError::InvalidKey(format!(
838 "Workspace key not found for: {}",
839 workspace_id
840 ))),
841 }
842 }
843 }
844 }
845
846 fn format_backup_string(&self, encrypted_key: &str) -> String {
847 let chars: Vec<char> = encrypted_key.chars().collect();
849 let mut result = String::new();
850
851 for (i, &ch) in chars.iter().enumerate() {
852 if i > 0 && i % 6 == 0 && i < chars.len() - 1 {
853 result.push('-');
854 }
855 result.push(ch);
856 }
857
858 if result.len() > 59 {
860 result.truncate(59);
862 }
863
864 result
865 }
866
867 fn parse_backup_string(&self, backup_string: &str) -> Result<String> {
868 Ok(backup_string.replace("-", ""))
870 }
871}
872
873impl Default for WorkspaceKeyManager {
874 fn default() -> Self {
875 Self::new()
876 }
877}
878
879#[derive(Debug, Clone, Serialize, Deserialize)]
881pub struct AutoEncryptionConfig {
882 pub enabled: bool,
884 pub sensitive_headers: Vec<String>,
886 pub sensitive_fields: Vec<String>,
888 pub sensitive_env_vars: Vec<String>,
890 pub sensitive_patterns: Vec<String>,
892}
893
894impl Default for AutoEncryptionConfig {
895 fn default() -> Self {
896 Self {
897 enabled: false,
898 sensitive_headers: vec![
899 "authorization".to_string(),
900 "x-api-key".to_string(),
901 "x-auth-token".to_string(),
902 "cookie".to_string(),
903 "set-cookie".to_string(),
904 ],
905 sensitive_fields: vec![
906 "password".to_string(),
907 "token".to_string(),
908 "secret".to_string(),
909 "key".to_string(),
910 "credentials".to_string(),
911 ],
912 sensitive_env_vars: vec![
913 "API_KEY".to_string(),
914 "SECRET_KEY".to_string(),
915 "PASSWORD".to_string(),
916 "TOKEN".to_string(),
917 "DATABASE_URL".to_string(),
918 ],
919 sensitive_patterns: vec![
920 r"\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b".to_string(), r"\b\d{3}[-\s]?\d{2}[-\s]?\d{4}\b".to_string(), ],
923 }
924 }
925}
926
927pub struct AutoEncryptionProcessor {
932 config: AutoEncryptionConfig,
934 workspace_manager: WorkspaceKeyManager,
936 workspace_id: String,
938}
939
940impl AutoEncryptionProcessor {
941 pub fn new(workspace_id: &str, config: AutoEncryptionConfig) -> Self {
943 Self {
944 config,
945 workspace_manager: WorkspaceKeyManager::new(),
946 workspace_id: workspace_id.to_string(),
947 }
948 }
949
950 pub fn process_headers(
952 &self,
953 headers: &mut std::collections::HashMap<String, String>,
954 ) -> Result<()> {
955 if !self.config.enabled {
956 return Ok(());
957 }
958
959 let workspace_key = self.workspace_manager.get_workspace_key(&self.workspace_id)?;
960
961 for (key, value) in headers.iter_mut() {
962 if self.is_sensitive_header(key) && !self.is_already_encrypted(value) {
963 *value = workspace_key.encrypt_chacha20(value, Some(key.as_bytes()))?;
964 }
965 }
966
967 Ok(())
968 }
969
970 pub fn process_json(&self, json: &mut serde_json::Value) -> Result<()> {
972 if !self.config.enabled {
973 return Ok(());
974 }
975
976 let workspace_key = self.workspace_manager.get_workspace_key(&self.workspace_id)?;
977 self.process_json_recursive(json, &workspace_key, Vec::new())?;
978
979 Ok(())
980 }
981
982 pub fn process_env_vars(
984 &self,
985 env_vars: &mut std::collections::HashMap<String, String>,
986 ) -> Result<()> {
987 if !self.config.enabled {
988 return Ok(());
989 }
990
991 let workspace_key = self.workspace_manager.get_workspace_key(&self.workspace_id)?;
992
993 for (key, value) in env_vars.iter_mut() {
994 if self.is_sensitive_env_var(key) && !self.is_already_encrypted(value) {
995 *value = workspace_key.encrypt_chacha20(value, Some(key.as_bytes()))?;
996 }
997 }
998
999 Ok(())
1000 }
1001
1002 fn is_sensitive_header(&self, header_name: &str) -> bool {
1004 self.config
1005 .sensitive_headers
1006 .iter()
1007 .any(|h| h.eq_ignore_ascii_case(header_name))
1008 }
1009
1010 fn is_sensitive_env_var(&self, var_name: &str) -> bool {
1012 self.config.sensitive_env_vars.iter().any(|v| v.eq_ignore_ascii_case(var_name))
1013 }
1014
1015 fn is_sensitive_field(&self, field_path: &[String]) -> bool {
1017 let default_field = String::new();
1018 let field_name = field_path.last().unwrap_or(&default_field);
1019
1020 if self.config.sensitive_fields.iter().any(|f| f.eq_ignore_ascii_case(field_name)) {
1022 return true;
1023 }
1024
1025 let path_str = field_path.join(".");
1027 for pattern in &self.config.sensitive_patterns {
1028 if let Ok(regex) = regex::Regex::new(pattern) {
1029 if regex.is_match(&path_str) || regex.is_match(field_name) {
1030 return true;
1031 }
1032 }
1033 }
1034
1035 false
1036 }
1037
1038 fn is_already_encrypted(&self, value: &str) -> bool {
1040 value.len() > 100 && general_purpose::STANDARD.decode(value).is_ok()
1042 }
1043
1044 fn process_json_recursive(
1046 &self,
1047 json: &mut serde_json::Value,
1048 workspace_key: &EncryptionKey,
1049 current_path: Vec<String>,
1050 ) -> Result<()> {
1051 match json {
1052 serde_json::Value::Object(obj) => {
1053 for (key, value) in obj.iter_mut() {
1054 let mut new_path = current_path.clone();
1055 new_path.push(key.clone());
1056
1057 if let serde_json::Value::String(ref mut s) = value {
1058 if self.is_sensitive_field(&new_path) && !self.is_already_encrypted(s) {
1059 let path_str = new_path.join(".");
1060 let path_bytes = path_str.as_bytes();
1061 *s = workspace_key.encrypt_chacha20(s, Some(path_bytes))?;
1062 }
1063 } else {
1064 self.process_json_recursive(value, workspace_key, new_path)?;
1065 }
1066 }
1067 }
1068 serde_json::Value::Array(arr) => {
1069 for (index, item) in arr.iter_mut().enumerate() {
1070 let mut new_path = current_path.clone();
1071 new_path.push(index.to_string());
1072 self.process_json_recursive(item, workspace_key, new_path)?;
1073 }
1074 }
1075 _ => {} }
1077
1078 Ok(())
1079 }
1080}
1081
1082pub mod utils {
1084 use super::*;
1085
1086 pub async fn is_encryption_enabled_for_workspace(
1088 persistence: &WorkspacePersistence,
1089 workspace_id: &str,
1090 ) -> Result<bool> {
1091 if let Ok(workspace) = persistence.load_workspace(workspace_id).await {
1093 return Ok(workspace.config.auto_encryption.enabled);
1094 }
1095 let manager = WorkspaceKeyManager::new();
1097 Ok(manager.has_workspace_key(workspace_id))
1098 }
1099
1100 pub async fn get_auto_encryption_config(
1102 persistence: &WorkspacePersistence,
1103 workspace_id: &str,
1104 ) -> Result<AutoEncryptionConfig> {
1105 let workspace = persistence.load_workspace(workspace_id).await.map_err(|e| {
1106 EncryptionError::Generic {
1107 message: format!("Failed to load workspace: {}", e),
1108 }
1109 })?;
1110 Ok(workspace.config.auto_encryption)
1111 }
1112
1113 pub fn encrypt_for_workspace(workspace_id: &str, data: &str) -> Result<String> {
1115 let manager = WorkspaceKeyManager::new();
1116 let key = manager.get_workspace_key(workspace_id)?;
1117 key.encrypt_chacha20(data, None)
1118 }
1119
1120 pub fn decrypt_for_workspace(workspace_id: &str, encrypted_data: &str) -> Result<String> {
1122 let manager = WorkspaceKeyManager::new();
1123 let key = manager.get_workspace_key(workspace_id)?;
1124 key.decrypt_chacha20(encrypted_data, None)
1125 }
1126}
1127
1128pub fn encrypt_with_key(
1141 key_id: &str,
1142 plaintext: &str,
1143 associated_data: Option<&[u8]>,
1144) -> Result<String> {
1145 let store = get_key_store()
1146 .ok_or_else(|| EncryptionError::InvalidKey("Key store not initialized".to_string()))?;
1147
1148 let key = store
1149 .get_key(key_id)
1150 .ok_or_else(|| EncryptionError::InvalidKey(format!("Key '{}' not found", key_id)))?;
1151
1152 key.encrypt(plaintext, associated_data)
1153}
1154
1155pub fn decrypt_with_key(
1168 key_id: &str,
1169 ciphertext: &str,
1170 associated_data: Option<&[u8]>,
1171) -> Result<String> {
1172 let store = get_key_store()
1173 .ok_or_else(|| EncryptionError::InvalidKey("Key store not initialized".to_string()))?;
1174
1175 let key = store
1176 .get_key(key_id)
1177 .ok_or_else(|| EncryptionError::InvalidKey(format!("Key '{}' not found", key_id)))?;
1178
1179 key.decrypt(ciphertext, associated_data)
1180}
1181
1182mod algorithms;
1183mod auto_encryption;
1184mod derivation;
1185mod errors;
1186mod key_management;
1187mod key_rotation;
1188
1189#[cfg(test)]
1190mod tests {
1191 use super::*;
1192 use once_cell::sync::Lazy;
1193 use std::sync::Mutex;
1194
1195 static MASTER_KEY_TEST_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
1196
1197 #[test]
1198 fn test_aes_gcm_encrypt_decrypt() {
1199 let key = EncryptionKey::from_password_pbkdf2(
1200 "test_password",
1201 None,
1202 EncryptionAlgorithm::Aes256Gcm,
1203 )
1204 .unwrap();
1205
1206 let plaintext = "Hello, World!";
1207 let ciphertext = key.encrypt(plaintext, None).unwrap();
1208 let decrypted = key.decrypt(&ciphertext, None).unwrap();
1209
1210 assert_eq!(plaintext, decrypted);
1211 }
1212
1213 #[test]
1214 fn test_chacha20_encrypt_decrypt() {
1215 let key = EncryptionKey::from_password_pbkdf2(
1216 "test_password",
1217 None,
1218 EncryptionAlgorithm::ChaCha20Poly1305,
1219 )
1220 .unwrap();
1221
1222 let plaintext = "Hello, World!";
1223 let ciphertext = key.encrypt(plaintext, None).unwrap();
1224 let decrypted = key.decrypt(&ciphertext, None).unwrap();
1225
1226 assert_eq!(plaintext, decrypted);
1227 }
1228
1229 #[test]
1230 fn test_key_store() {
1231 let mut store = KeyStore::new();
1232
1233 store
1234 .derive_and_store_key(
1235 "test_key".to_string(),
1236 "test_password",
1237 EncryptionAlgorithm::Aes256Gcm,
1238 KeyDerivationMethod::Pbkdf2,
1239 )
1240 .unwrap();
1241
1242 assert!(store.get_key("test_key").is_some());
1243 assert!(store.list_keys().contains(&"test_key".to_string()));
1244
1245 store.remove_key("test_key");
1246 assert!(store.get_key("test_key").is_none());
1247 }
1248
1249 #[test]
1250 fn test_invalid_key_length() {
1251 let result = EncryptionKey::new(EncryptionAlgorithm::Aes256Gcm, vec![1, 2, 3]);
1252 assert!(matches!(result, Err(EncryptionError::InvalidKey(_))));
1253 }
1254
1255 #[test]
1256 fn test_invalid_ciphertext() {
1257 let key = EncryptionKey::from_password_pbkdf2("test", None, EncryptionAlgorithm::Aes256Gcm)
1258 .unwrap();
1259 let result = key.decrypt("invalid_base64!", None);
1260 assert!(matches!(result, Err(EncryptionError::InvalidCiphertext(_))));
1261 }
1262
1263 #[test]
1264 fn test_chacha20_encrypt_decrypt_12byte_nonce() {
1265 let key = EncryptionKey::from_password_pbkdf2(
1266 "test_password",
1267 None,
1268 EncryptionAlgorithm::ChaCha20Poly1305,
1269 )
1270 .unwrap();
1271
1272 let plaintext = "Hello, World! This is a test of ChaCha20-Poly1305 with 12-byte nonce.";
1273 let ciphertext = key.encrypt_chacha20(plaintext, None).unwrap();
1274 let decrypted = key.decrypt_chacha20(&ciphertext, None).unwrap();
1275
1276 assert_eq!(plaintext, decrypted);
1277 }
1278
1279 #[test]
1280 fn test_secure_function_template() {
1281 use crate::templating::expand_str;
1282
1283 let template = r#"{{secure "test message"}}"#;
1285 let result = expand_str(template);
1286
1287 assert_ne!(result, "test message");
1289 assert!(!result.is_empty());
1290
1291 assert!(general_purpose::STANDARD.decode(&result).is_ok());
1293 }
1294
1295 #[test]
1296 fn test_master_key_manager() {
1297 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1298
1299 let manager = MasterKeyManager::new();
1300
1301 #[cfg(any(target_os = "macos", target_os = "linux"))]
1303 {
1304 if let Ok(home) = std::env::var("HOME") {
1305 let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
1306 let _ = std::fs::remove_file(&key_path);
1307 }
1308 }
1309
1310 assert!(!manager.has_master_key());
1312
1313 manager.generate_master_key().unwrap();
1315 assert!(manager.has_master_key());
1316
1317 let master_key = manager.get_master_key().unwrap();
1319 assert_eq!(master_key.algorithm, EncryptionAlgorithm::ChaCha20Poly1305);
1320 }
1321
1322 #[test]
1323 fn test_workspace_key_manager() {
1324 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1325
1326 let master_manager = MasterKeyManager::new();
1328 let needs_generation =
1329 !master_manager.has_master_key() || master_manager.get_master_key().is_err();
1330 if needs_generation && master_manager.generate_master_key().is_err() {
1331 return;
1333 }
1334
1335 let workspace_manager = WorkspaceKeyManager::new();
1336 let workspace_id = &format!("test_workspace_{}", uuid::Uuid::new_v4());
1337
1338 assert!(!workspace_manager.has_workspace_key(workspace_id));
1340
1341 let encrypted_key = workspace_manager.generate_workspace_key(workspace_id).unwrap();
1343 assert!(workspace_manager.has_workspace_key(workspace_id));
1344 assert!(!encrypted_key.is_empty());
1345
1346 let workspace_key = workspace_manager.get_workspace_key(workspace_id).unwrap();
1348 assert_eq!(workspace_key.algorithm, EncryptionAlgorithm::ChaCha20Poly1305);
1349
1350 let test_data = "sensitive workspace data";
1352 let ciphertext = workspace_key.encrypt_chacha20(test_data, None).unwrap();
1353 let decrypted = workspace_key.decrypt_chacha20(&ciphertext, None).unwrap();
1354 assert_eq!(test_data, decrypted);
1355 }
1356
1357 #[test]
1358 fn test_backup_string_formatting() {
1359 let manager = WorkspaceKeyManager::new();
1360
1361 let test_key = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrst";
1363 let backup = manager.format_backup_string(test_key);
1364
1365 assert!(backup.contains('-'));
1367
1368 let parsed = manager.parse_backup_string(&backup).unwrap();
1370 assert_eq!(parsed, test_key);
1371 }
1372
1373 #[test]
1374 fn test_auto_encryption_processor() {
1375 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1376
1377 let master_manager = MasterKeyManager::new();
1379 let needs_key =
1380 !master_manager.has_master_key() || master_manager.get_master_key().is_err();
1381 if needs_key && master_manager.generate_master_key().is_err() {
1382 eprintln!("Skipping test_auto_encryption_processor: Failed to generate master key");
1383 return;
1384 }
1385
1386 let workspace_manager = WorkspaceKeyManager::new();
1387 let workspace_id = &format!("test_auto_encrypt_workspace_{}", uuid::Uuid::new_v4());
1388
1389 workspace_manager.generate_workspace_key(workspace_id).unwrap();
1390
1391 let config = AutoEncryptionConfig {
1392 enabled: true,
1393 ..AutoEncryptionConfig::default()
1394 };
1395
1396 let processor = AutoEncryptionProcessor::new(workspace_id, config);
1397
1398 let mut headers = std::collections::HashMap::new();
1400 headers.insert("Authorization".to_string(), "Bearer my-secret-token".to_string());
1401 headers.insert("Content-Type".to_string(), "application/json".to_string());
1402
1403 processor.process_headers(&mut headers).unwrap();
1404
1405 assert_ne!(headers["Authorization"], "Bearer my-secret-token");
1407 assert!(general_purpose::STANDARD.decode(&headers["Authorization"]).is_ok());
1408
1409 assert_eq!(headers["Content-Type"], "application/json");
1411 }
1412
1413 #[test]
1414 #[ignore] fn test_json_field_encryption() {
1416 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1417
1418 let master_manager = MasterKeyManager::new();
1420 if !master_manager.has_master_key() || master_manager.get_master_key().is_err() {
1421 master_manager.generate_master_key().unwrap();
1422 }
1423
1424 let workspace_manager = WorkspaceKeyManager::new();
1425 let workspace_id = &format!("test_json_workspace_{}", uuid::Uuid::new_v4());
1426
1427 workspace_manager.generate_workspace_key(workspace_id).unwrap();
1428
1429 let another_manager = WorkspaceKeyManager::new();
1431 assert!(another_manager.has_workspace_key(workspace_id));
1432
1433 let config = AutoEncryptionConfig {
1434 enabled: true,
1435 ..AutoEncryptionConfig::default()
1436 };
1437
1438 let processor = AutoEncryptionProcessor::new(workspace_id, config);
1439
1440 let mut json = serde_json::json!({
1442 "username": "testuser",
1443 "password": "secret123",
1444 "email": "test@example.com",
1445 "nested": {
1446 "token": "my-api-token",
1447 "normal_field": "normal_value"
1448 }
1449 });
1450
1451 processor.process_json(&mut json).unwrap();
1452
1453 assert_ne!(json["password"], "secret123");
1455 assert_ne!(json["nested"]["token"], "my-api-token");
1456
1457 assert_eq!(json["username"], "testuser");
1459 assert_eq!(json["email"], "test@example.com");
1460 assert_eq!(json["nested"]["normal_field"], "normal_value");
1461 }
1462
1463 #[test]
1464 fn test_env_var_encryption() {
1465 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1466
1467 let master_manager = MasterKeyManager::new();
1469 if !master_manager.has_master_key() || master_manager.get_master_key().is_err() {
1470 master_manager.generate_master_key().unwrap();
1471 }
1472
1473 let workspace_manager = WorkspaceKeyManager::new();
1474 let workspace_id = &format!("test_env_workspace_{}", uuid::Uuid::new_v4());
1475
1476 workspace_manager.generate_workspace_key(workspace_id).unwrap();
1477
1478 let config = AutoEncryptionConfig {
1479 enabled: true,
1480 ..AutoEncryptionConfig::default()
1481 };
1482
1483 let processor = AutoEncryptionProcessor::new(workspace_id, config);
1484
1485 let mut env_vars = std::collections::HashMap::new();
1487 env_vars.insert("API_KEY".to_string(), "sk-1234567890abcdef".to_string());
1488 env_vars
1489 .insert("DATABASE_URL".to_string(), "postgres://user:pass@host:5432/db".to_string());
1490 env_vars.insert("NORMAL_VAR".to_string(), "normal_value".to_string());
1491
1492 processor.process_env_vars(&mut env_vars).unwrap();
1493
1494 assert_ne!(env_vars["API_KEY"], "sk-1234567890abcdef");
1496 assert_ne!(env_vars["DATABASE_URL"], "postgres://user:pass@host:5432/db");
1497
1498 assert_eq!(env_vars["NORMAL_VAR"], "normal_value");
1500 }
1501
1502 #[test]
1503 fn test_encryption_utils() {
1504 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1505
1506 let master_manager = MasterKeyManager::new();
1508 let needs_key =
1509 !master_manager.has_master_key() || master_manager.get_master_key().is_err();
1510 if needs_key && master_manager.generate_master_key().is_err() {
1511 eprintln!("Skipping test_encryption_utils: Failed to generate master key");
1512 return;
1513 }
1514
1515 let workspace_manager = WorkspaceKeyManager::new();
1516 let workspace_id = &format!("test_utils_workspace_{}", uuid::Uuid::new_v4());
1517 workspace_manager.generate_workspace_key(workspace_id).unwrap();
1518
1519 assert!(workspace_manager.has_workspace_key(workspace_id));
1521
1522 let another_manager = WorkspaceKeyManager::new();
1524 assert!(another_manager.has_workspace_key(workspace_id));
1525
1526 let test_data = "test data for utils";
1527 let encrypted = utils::encrypt_for_workspace(workspace_id, test_data).unwrap();
1528 let decrypted = utils::decrypt_for_workspace(workspace_id, &encrypted).unwrap();
1529
1530 assert_eq!(test_data, decrypted);
1531 assert_ne!(encrypted, test_data);
1532 }
1533}