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 tracing;
35
36pub type Result<T> = EncryptionResult<T>;
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55pub enum EncryptionAlgorithm {
56 Aes256Gcm,
58 ChaCha20Poly1305,
60}
61
62impl fmt::Display for EncryptionAlgorithm {
63 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64 match self {
65 EncryptionAlgorithm::Aes256Gcm => write!(f, "aes256-gcm"),
66 EncryptionAlgorithm::ChaCha20Poly1305 => write!(f, "chacha20-poly1305"),
67 }
68 }
69}
70
71#[derive(Debug, Clone, Copy, PartialEq, Eq)]
73pub enum KeyDerivationMethod {
74 Pbkdf2,
76 Argon2,
78}
79
80impl fmt::Display for KeyDerivationMethod {
81 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
82 match self {
83 KeyDerivationMethod::Pbkdf2 => write!(f, "pbkdf2"),
84 KeyDerivationMethod::Argon2 => write!(f, "argon2"),
85 }
86 }
87}
88
89pub struct EncryptionKey {
91 algorithm: EncryptionAlgorithm,
92 key_data: Vec<u8>,
93}
94
95impl EncryptionKey {
96 pub fn new(algorithm: EncryptionAlgorithm, key_data: Vec<u8>) -> Result<Self> {
98 let expected_len = match algorithm {
99 EncryptionAlgorithm::Aes256Gcm => 32, EncryptionAlgorithm::ChaCha20Poly1305 => 32, };
102
103 if key_data.len() != expected_len {
104 return Err(EncryptionError::InvalidKey {
105 message: format!(
106 "Key must be {} bytes for {}, got {}",
107 expected_len,
108 algorithm,
109 key_data.len()
110 ),
111 });
112 }
113
114 Ok(Self {
115 algorithm,
116 key_data,
117 })
118 }
119
120 pub fn from_password_pbkdf2(
122 password: &str,
123 salt: Option<&[u8]>,
124 algorithm: EncryptionAlgorithm,
125 ) -> Result<Self> {
126 let salt = salt
127 .map(|s| s.to_vec())
128 .unwrap_or_else(|| thread_rng().random::<[u8; 32]>().to_vec());
129
130 let mut key = vec![0u8; 32];
131 pbkdf2_hmac::<Sha256>(password.as_bytes(), &salt, 100_000, &mut key);
132
133 Self::new(algorithm, key)
134 }
135
136 pub fn from_password_argon2(
138 password: &str,
139 salt: Option<&[u8]>,
140 algorithm: EncryptionAlgorithm,
141 ) -> Result<Self> {
142 let salt_string = if let Some(salt) = salt {
143 SaltString::encode_b64(salt).map_err(|e| EncryptionError::KeyDerivationFailed {
144 message: e.to_string(),
145 })?
146 } else {
147 let mut salt_bytes = [0u8; 32];
149 thread_rng().fill(&mut salt_bytes);
150 SaltString::encode_b64(&salt_bytes).map_err(|e| {
151 EncryptionError::KeyDerivationFailed {
152 message: e.to_string(),
153 }
154 })?
155 };
156
157 let argon2 = Argon2::new(
158 argon2::Algorithm::Argon2id,
159 argon2::Version::V0x13,
160 Params::new(65536, 3, 1, Some(32)).map_err(|e| {
161 EncryptionError::KeyDerivationFailed {
162 message: e.to_string(),
163 }
164 })?,
165 );
166
167 let hash = argon2.hash_password(password.as_bytes(), &salt_string).map_err(|e| {
168 EncryptionError::KeyDerivationFailed {
169 message: e.to_string(),
170 }
171 })?;
172
173 let key_bytes = hash
174 .hash
175 .ok_or_else(|| EncryptionError::KeyDerivationFailed {
176 message: "Failed to derive key hash".to_string(),
177 })?
178 .as_bytes()
179 .to_vec();
180 Self::new(algorithm, key_bytes)
181 }
182
183 pub fn encrypt(&self, plaintext: &str, associated_data: Option<&[u8]>) -> Result<String> {
185 match self.algorithm {
186 EncryptionAlgorithm::Aes256Gcm => self.encrypt_aes_gcm(plaintext, associated_data),
187 EncryptionAlgorithm::ChaCha20Poly1305 => {
188 self.encrypt_chacha20(plaintext, associated_data)
189 }
190 }
191 }
192
193 pub fn decrypt(&self, ciphertext: &str, associated_data: Option<&[u8]>) -> Result<String> {
195 match self.algorithm {
196 EncryptionAlgorithm::Aes256Gcm => self.decrypt_aes_gcm(ciphertext, associated_data),
197 EncryptionAlgorithm::ChaCha20Poly1305 => {
198 self.decrypt_chacha20(ciphertext, associated_data)
199 }
200 }
201 }
202
203 fn encrypt_aes_gcm(&self, plaintext: &str, associated_data: Option<&[u8]>) -> Result<String> {
204 let key_array: [u8; 32] =
206 self.key_data.as_slice().try_into().map_err(|_| EncryptionError::InvalidKey {
207 message: "Key length must be 32 bytes".to_string(),
208 })?;
209 let cipher = Aes256Gcm::new(&key_array.into());
210 let nonce: [u8; 12] = thread_rng().random(); let nonce = Nonce::from(nonce);
212
213 let ciphertext = cipher.encrypt(&nonce, plaintext.as_bytes()).map_err(|e| {
214 EncryptionError::CipherOperationFailed {
215 message: e.to_string(),
216 }
217 })?;
218
219 let mut result = nonce.to_vec();
220 result.extend_from_slice(&ciphertext);
221
222 if let Some(aad) = associated_data {
224 result.extend_from_slice(aad);
225 }
226
227 Ok(general_purpose::STANDARD.encode(&result))
228 }
229
230 fn decrypt_aes_gcm(&self, ciphertext: &str, associated_data: Option<&[u8]>) -> Result<String> {
231 let data = general_purpose::STANDARD.decode(ciphertext).map_err(|e| {
232 EncryptionError::InvalidCiphertext {
233 message: e.to_string(),
234 }
235 })?;
236
237 if data.len() < 12 {
238 return Err(EncryptionError::InvalidCiphertext {
239 message: "Ciphertext too short".to_string(),
240 });
241 }
242
243 let nonce_array: [u8; 12] =
245 data[0..12].try_into().map_err(|_| EncryptionError::InvalidCiphertext {
246 message: "Nonce must be 12 bytes".to_string(),
247 })?;
248 let nonce = Nonce::from(nonce_array);
249
250 let ciphertext_len = if let Some(aad) = &associated_data {
251 let aad_len = aad.len();
253 data.len() - 12 - aad_len
254 } else {
255 data.len() - 12
256 };
257
258 let ciphertext = &data[12..12 + ciphertext_len];
259
260 let key_array: [u8; 32] =
262 self.key_data.as_slice().try_into().map_err(|_| EncryptionError::InvalidKey {
263 message: "Key length must be 32 bytes".to_string(),
264 })?;
265 let cipher = Aes256Gcm::new(&key_array.into());
266
267 let plaintext = cipher.decrypt(&nonce, ciphertext.as_ref()).map_err(|e| {
268 EncryptionError::CipherOperationFailed {
269 message: e.to_string(),
270 }
271 })?;
272
273 String::from_utf8(plaintext).map_err(|e| EncryptionError::CipherOperationFailed {
274 message: format!("Invalid UTF-8: {}", e),
275 })
276 }
277
278 pub fn encrypt_chacha20(
287 &self,
288 plaintext: &str,
289 _associated_data: Option<&[u8]>,
290 ) -> Result<String> {
291 let key = ChaChaKey::from_slice(&self.key_data);
292 let cipher = ChaCha20Poly1305::new(key);
293 let nonce: [u8; 12] = thread_rng().random(); let nonce = chacha20poly1305::Nonce::from_slice(&nonce);
295
296 let ciphertext = cipher.encrypt(nonce, plaintext.as_bytes()).map_err(|e| {
297 EncryptionError::CipherOperationFailed {
298 message: e.to_string(),
299 }
300 })?;
301
302 let mut result = nonce.to_vec();
303 result.extend_from_slice(&ciphertext);
304
305 Ok(general_purpose::STANDARD.encode(&result))
306 }
307
308 pub fn decrypt_chacha20(
317 &self,
318 ciphertext: &str,
319 _associated_data: Option<&[u8]>,
320 ) -> Result<String> {
321 let data = general_purpose::STANDARD.decode(ciphertext).map_err(|e| {
322 EncryptionError::InvalidCiphertext {
323 message: e.to_string(),
324 }
325 })?;
326
327 if data.len() < 12 {
328 return Err(EncryptionError::InvalidCiphertext {
329 message: "Ciphertext too short".to_string(),
330 });
331 }
332
333 let nonce = chacha20poly1305::Nonce::from_slice(&data[0..12]);
334 let ciphertext_data = &data[12..];
335 let key = ChaChaKey::from_slice(&self.key_data);
336 let cipher = ChaCha20Poly1305::new(key);
337
338 let plaintext = cipher.decrypt(nonce, ciphertext_data.as_ref()).map_err(|e| {
339 EncryptionError::CipherOperationFailed {
340 message: e.to_string(),
341 }
342 })?;
343
344 String::from_utf8(plaintext).map_err(|e| EncryptionError::CipherOperationFailed {
345 message: format!("Invalid UTF-8: {}", e),
346 })
347 }
348}
349
350pub struct KeyStore {
352 keys: std::collections::HashMap<String, EncryptionKey>,
353}
354
355impl KeyStore {
356 pub fn new() -> Self {
358 Self {
359 keys: std::collections::HashMap::new(),
360 }
361 }
362
363 pub fn store_key(&mut self, id: String, key: EncryptionKey) {
365 self.keys.insert(id, key);
366 }
367
368 pub fn get_key(&self, id: &str) -> Option<&EncryptionKey> {
370 self.keys.get(id)
371 }
372
373 pub fn remove_key(&mut self, id: &str) -> bool {
375 self.keys.remove(id).is_some()
376 }
377
378 pub fn list_keys(&self) -> Vec<String> {
380 self.keys.keys().cloned().collect()
381 }
382
383 pub fn derive_and_store_key(
385 &mut self,
386 id: String,
387 password: &str,
388 algorithm: EncryptionAlgorithm,
389 method: KeyDerivationMethod,
390 ) -> Result<()> {
391 let key = match method {
392 KeyDerivationMethod::Pbkdf2 => {
393 EncryptionKey::from_password_pbkdf2(password, None, algorithm)?
394 }
395 KeyDerivationMethod::Argon2 => {
396 EncryptionKey::from_password_argon2(password, None, algorithm)?
397 }
398 };
399 self.store_key(id, key);
400 Ok(())
401 }
402}
403
404impl Default for KeyStore {
405 fn default() -> Self {
406 Self::new()
407 }
408}
409
410static KEY_STORE: once_cell::sync::OnceCell<KeyStore> = once_cell::sync::OnceCell::new();
412
413pub fn init_key_store() -> &'static KeyStore {
422 KEY_STORE.get_or_init(KeyStore::default)
423}
424
425pub fn get_key_store() -> Option<&'static KeyStore> {
432 KEY_STORE.get()
433}
434
435pub struct MasterKeyManager {
442 _service_name: String,
444 _account_name: String,
446}
447
448impl MasterKeyManager {
449 pub fn new() -> Self {
451 Self {
452 _service_name: "com.mockforge.encryption".to_string(),
453 _account_name: "master_key".to_string(),
454 }
455 }
456
457 pub fn generate_master_key(&self) -> Result<()> {
459 let master_key_bytes: [u8; 32] = rand::random();
460 let master_key_b64 = general_purpose::STANDARD.encode(master_key_bytes);
461
462 #[cfg(target_os = "macos")]
465 {
466 self.store_in_macos_keychain(&master_key_b64)?;
468 }
469 #[cfg(target_os = "linux")]
470 {
471 self.store_in_linux_keyring(&master_key_b64)?;
473 }
474 #[cfg(target_os = "windows")]
475 {
476 self.store_in_windows_credential_manager(&master_key_b64)?;
478 }
479
480 Ok(())
481 }
482
483 pub fn get_master_key(&self) -> Result<EncryptionKey> {
485 let master_key_b64 = self.retrieve_from_keychain()?;
486 let master_key_bytes = general_purpose::STANDARD.decode(master_key_b64).map_err(|e| {
487 EncryptionError::InvalidKey {
488 message: e.to_string(),
489 }
490 })?;
491
492 if master_key_bytes.len() != 32 {
493 return Err(EncryptionError::InvalidKey {
494 message: "Invalid master key length".to_string(),
495 });
496 }
497
498 EncryptionKey::new(EncryptionAlgorithm::ChaCha20Poly1305, master_key_bytes)
499 }
500
501 pub fn has_master_key(&self) -> bool {
503 self.retrieve_from_keychain().is_ok()
504 }
505
506 #[cfg(target_os = "macos")]
508 fn store_in_macos_keychain(&self, key: &str) -> Result<()> {
509 use std::os::unix::fs::PermissionsExt;
510
511 let home = std::env::var("HOME").map_err(|_| EncryptionError::InvalidKey {
512 message: "HOME environment variable not set".to_string(),
513 })?;
514 let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
515
516 if let Some(parent) = key_path.parent() {
518 std::fs::create_dir_all(parent).map_err(|e| EncryptionError::InvalidKey {
519 message: format!("Failed to create directory: {}", e),
520 })?;
521 }
522
523 std::fs::write(&key_path, key).map_err(|e| EncryptionError::InvalidKey {
525 message: format!("Failed to write master key: {}", e),
526 })?;
527
528 let mut perms = std::fs::metadata(&key_path)
530 .map_err(|e| EncryptionError::InvalidKey {
531 message: format!("Failed to get metadata: {}", e),
532 })?
533 .permissions();
534 perms.set_mode(0o600);
535 std::fs::set_permissions(&key_path, perms).map_err(|e| EncryptionError::InvalidKey {
536 message: format!("Failed to set permissions: {}", e),
537 })?;
538
539 Ok(())
540 }
541
542 #[cfg(target_os = "linux")]
543 fn store_in_linux_keyring(&self, key: &str) -> Result<()> {
544 use std::os::unix::fs::PermissionsExt;
545
546 let home = std::env::var("HOME").map_err(|_| EncryptionError::InvalidKey {
547 message: "HOME environment variable not set".to_string(),
548 })?;
549 let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
550
551 if let Some(parent) = key_path.parent() {
553 std::fs::create_dir_all(parent).map_err(|e| EncryptionError::InvalidKey {
554 message: format!("Failed to create directory: {}", e),
555 })?;
556 }
557
558 std::fs::write(&key_path, key).map_err(|e| EncryptionError::InvalidKey {
560 message: format!("Failed to write master key: {}", e),
561 })?;
562
563 let mut perms = std::fs::metadata(&key_path)
565 .map_err(|e| EncryptionError::InvalidKey {
566 message: format!("Failed to get metadata: {}", e),
567 })?
568 .permissions();
569 perms.set_mode(0o600);
570 std::fs::set_permissions(&key_path, perms).map_err(|e| EncryptionError::InvalidKey {
571 message: format!("Failed to set permissions: {}", e),
572 })?;
573
574 Ok(())
575 }
576
577 #[cfg(target_os = "windows")]
578 fn store_in_windows_credential_manager(&self, key: &str) -> Result<()> {
579 use std::ffi::OsStr;
580 use std::os::windows::ffi::OsStrExt;
581 use windows::core::PWSTR;
582 use windows::Win32::Security::Credentials::{
583 CredWriteW, CREDENTIALW, CRED_FLAGS, CRED_PERSIST_LOCAL_MACHINE, CRED_TYPE_GENERIC,
584 };
585
586 let target_name = "MockForge/MasterKey";
587 let mut target_name_wide: Vec<u16> =
588 OsStr::new(target_name).encode_wide().chain(std::iter::once(0)).collect();
589
590 let mut credential_blob: Vec<u16> =
591 OsStr::new(key).encode_wide().chain(std::iter::once(0)).collect();
592
593 let credential = CREDENTIALW {
594 Flags: CRED_FLAGS::default(),
595 Type: CRED_TYPE_GENERIC,
596 TargetName: PWSTR::from_raw(target_name_wide.as_mut_ptr()),
597 Comment: PWSTR::null(),
598 LastWritten: windows::Win32::Foundation::FILETIME::default(),
599 CredentialBlobSize: (credential_blob.len() * 2) as u32,
600 CredentialBlob: credential_blob.as_mut_ptr() as *mut u8,
601 Persist: CRED_PERSIST_LOCAL_MACHINE,
602 AttributeCount: 0,
603 Attributes: std::ptr::null_mut(),
604 TargetAlias: PWSTR::null(),
605 UserName: PWSTR::null(),
606 };
607
608 unsafe {
614 CredWriteW(&credential, 0).map_err(|e| EncryptionError::InvalidKey {
615 message: format!("Failed to store credential: {:?}", e),
616 })?;
617 }
618
619 Ok(())
620 }
621
622 fn retrieve_from_keychain(&self) -> Result<String> {
623 #[cfg(target_os = "macos")]
625 {
626 let home = std::env::var("HOME").map_err(|_| EncryptionError::InvalidKey {
627 message: "HOME environment variable not set".to_string(),
628 })?;
629 let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
630 let key_str =
631 std::fs::read_to_string(&key_path).map_err(|_| EncryptionError::InvalidKey {
632 message: "Master key not found in keychain".to_string(),
633 })?;
634 Ok(key_str.trim().to_string())
636 }
637
638 #[cfg(target_os = "linux")]
639 {
640 let home = std::env::var("HOME").map_err(|_| EncryptionError::InvalidKey {
641 message: "HOME environment variable not set".to_string(),
642 })?;
643 let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
644 let key_str =
645 std::fs::read_to_string(&key_path).map_err(|_| EncryptionError::InvalidKey {
646 message: "Master key not found in keychain".to_string(),
647 })?;
648 Ok(key_str.trim().to_string())
650 }
651
652 #[cfg(target_os = "windows")]
653 {
654 self.retrieve_from_windows_credential_manager()
656 }
657
658 #[cfg(not(any(target_os = "macos", target_os = "linux", target_os = "windows")))]
659 {
660 let key_str =
662 std::env::var("MOCKFORGE_MASTER_KEY").map_err(|_| EncryptionError::InvalidKey {
663 message: "Master key not found in keychain".to_string(),
664 })?;
665 Ok(key_str.trim().to_string())
667 }
668 }
669
670 #[cfg(target_os = "windows")]
671 fn retrieve_from_windows_credential_manager(&self) -> Result<String> {
672 use std::ffi::OsString;
673 use std::os::windows::ffi::{OsStrExt, OsStringExt};
674 use windows::core::PCWSTR;
675 use windows::Win32::Security::Credentials::{
676 CredFree, CredReadW, CREDENTIALW, CRED_TYPE_GENERIC,
677 };
678
679 let target_name = "MockForge/MasterKey";
680 let target_name_wide: Vec<u16> = std::ffi::OsStr::new(target_name)
681 .encode_wide()
682 .chain(std::iter::once(0))
683 .collect();
684
685 let mut credential_ptr: *mut CREDENTIALW = std::ptr::null_mut();
686
687 unsafe {
692 CredReadW(
693 PCWSTR::from_raw(target_name_wide.as_ptr()),
694 CRED_TYPE_GENERIC,
695 None,
696 &mut credential_ptr,
697 )
698 .map_err(|e| EncryptionError::InvalidKey {
699 message: format!("Failed to read credential: {:?}", e),
700 })?;
701
702 if credential_ptr.is_null() {
703 return Err(EncryptionError::InvalidKey {
704 message: "Credential not found".to_string(),
705 });
706 }
707
708 let credential = &*credential_ptr;
710
711 let blob_slice = std::slice::from_raw_parts(
714 credential.CredentialBlob as *const u16,
715 credential.CredentialBlobSize as usize / 2, );
717
718 let credential_str = OsString::from_wide(blob_slice)
719 .to_string_lossy()
720 .trim_end_matches('\0')
721 .to_string();
722
723 CredFree(credential_ptr as *const std::ffi::c_void);
725
726 Ok(credential_str)
727 }
728 }
729}
730
731impl Default for MasterKeyManager {
732 fn default() -> Self {
733 Self::new()
734 }
735}
736
737pub struct WorkspaceKeyManager {
742 master_key_manager: MasterKeyManager,
744 key_storage: std::cell::RefCell<FileKeyStorage>,
746}
747
748impl WorkspaceKeyManager {
749 pub fn new() -> Self {
751 Self {
752 master_key_manager: MasterKeyManager::new(),
753 key_storage: std::cell::RefCell::new(FileKeyStorage::new()),
754 }
755 }
756
757 pub fn with_storage_path<P: AsRef<std::path::Path>>(path: P) -> Self {
759 Self {
760 master_key_manager: MasterKeyManager::new(),
761 key_storage: std::cell::RefCell::new(FileKeyStorage::with_path(path)),
762 }
763 }
764
765 pub fn generate_workspace_key(&self, workspace_id: &str) -> Result<String> {
767 let workspace_key_bytes: [u8; 32] = rand::random();
769
770 let master_key = self.master_key_manager.get_master_key()?;
772
773 let workspace_key_b64 = master_key.encrypt_chacha20(
775 &general_purpose::STANDARD.encode(workspace_key_bytes),
776 Some(workspace_id.as_bytes()),
777 )?;
778
779 self.store_workspace_key(workspace_id, &workspace_key_b64)?;
781
782 Ok(workspace_key_b64)
783 }
784
785 pub fn get_workspace_key(&self, workspace_id: &str) -> Result<EncryptionKey> {
787 let encrypted_key_b64 = self.retrieve_workspace_key(workspace_id)?;
788 let master_key = self.master_key_manager.get_master_key()?;
789
790 let decrypted_key_b64 =
791 master_key.decrypt_chacha20(&encrypted_key_b64, Some(workspace_id.as_bytes()))?;
792
793 let workspace_key_bytes =
794 general_purpose::STANDARD.decode(decrypted_key_b64).map_err(|e| {
795 EncryptionError::InvalidKey {
796 message: e.to_string(),
797 }
798 })?;
799
800 if workspace_key_bytes.len() != 32 {
801 return Err(EncryptionError::InvalidKey {
802 message: "Invalid workspace key length".to_string(),
803 });
804 }
805
806 EncryptionKey::new(EncryptionAlgorithm::ChaCha20Poly1305, workspace_key_bytes)
807 }
808
809 pub fn has_workspace_key(&self, workspace_id: &str) -> bool {
811 self.retrieve_workspace_key(workspace_id).is_ok()
812 }
813
814 pub fn generate_workspace_key_backup(&self, workspace_id: &str) -> Result<String> {
816 let encrypted_key = self.retrieve_workspace_key(workspace_id)?;
817
818 let backup_string = self.format_backup_string(&encrypted_key);
821
822 Ok(backup_string)
823 }
824
825 pub fn restore_workspace_key_from_backup(
827 &self,
828 workspace_id: &str,
829 backup_string: &str,
830 ) -> Result<()> {
831 let encrypted_key = self.parse_backup_string(backup_string)?;
832 self.store_workspace_key(workspace_id, &encrypted_key)
833 }
834
835 fn store_workspace_key(&self, workspace_id: &str, encrypted_key: &str) -> Result<()> {
837 self.key_storage
838 .borrow_mut()
839 .store_key(&workspace_id.to_string(), encrypted_key.as_bytes())
840 .map_err(|e| EncryptionError::InvalidKey {
841 message: format!("Failed to store workspace key: {:?}", e),
842 })
843 }
844
845 fn retrieve_workspace_key(&self, workspace_id: &str) -> Result<String> {
846 match self.key_storage.borrow().retrieve_key(&workspace_id.to_string()) {
848 Ok(encrypted_bytes) => {
849 String::from_utf8(encrypted_bytes).map_err(|e| EncryptionError::InvalidKey {
850 message: format!("Invalid UTF-8 in stored key: {}", e),
851 })
852 }
853 Err(_) => {
854 let old_key_file = format!("workspace_{}_key.enc", workspace_id);
856 match std::fs::read_to_string(&old_key_file) {
857 Ok(encrypted_key) => {
858 if let Err(e) = self
860 .key_storage
861 .borrow_mut()
862 .store_key(&workspace_id.to_string(), encrypted_key.as_bytes())
863 {
864 tracing::warn!(
865 "Failed to migrate workspace key to new storage: {:?}",
866 e
867 );
868 } else {
869 let _ = std::fs::remove_file(&old_key_file);
871 }
872 Ok(encrypted_key)
873 }
874 Err(_) => Err(EncryptionError::InvalidKey {
875 message: format!("Workspace key not found for: {}", workspace_id),
876 }),
877 }
878 }
879 }
880 }
881
882 fn format_backup_string(&self, encrypted_key: &str) -> String {
883 let chars: Vec<char> = encrypted_key.chars().collect();
885 let mut result = String::new();
886
887 for (i, &ch) in chars.iter().enumerate() {
888 if i > 0 && i % 6 == 0 && i < chars.len() - 1 {
889 result.push('-');
890 }
891 result.push(ch);
892 }
893
894 if result.len() > 59 {
896 result.truncate(59);
898 }
899
900 result
901 }
902
903 fn parse_backup_string(&self, backup_string: &str) -> Result<String> {
904 Ok(backup_string.replace("-", ""))
906 }
907}
908
909impl Default for WorkspaceKeyManager {
910 fn default() -> Self {
911 Self::new()
912 }
913}
914
915#[derive(Debug, Clone, Serialize, Deserialize)]
917pub struct AutoEncryptionConfig {
918 pub enabled: bool,
920 pub sensitive_headers: Vec<String>,
922 pub sensitive_fields: Vec<String>,
924 pub sensitive_env_vars: Vec<String>,
926 pub sensitive_patterns: Vec<String>,
928}
929
930impl Default for AutoEncryptionConfig {
931 fn default() -> Self {
932 Self {
933 enabled: false,
934 sensitive_headers: vec![
935 "authorization".to_string(),
936 "x-api-key".to_string(),
937 "x-auth-token".to_string(),
938 "cookie".to_string(),
939 "set-cookie".to_string(),
940 ],
941 sensitive_fields: vec![
942 "password".to_string(),
943 "token".to_string(),
944 "secret".to_string(),
945 "key".to_string(),
946 "credentials".to_string(),
947 ],
948 sensitive_env_vars: vec![
949 "API_KEY".to_string(),
950 "SECRET_KEY".to_string(),
951 "PASSWORD".to_string(),
952 "TOKEN".to_string(),
953 "DATABASE_URL".to_string(),
954 ],
955 sensitive_patterns: vec![
956 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(), ],
959 }
960 }
961}
962
963pub struct AutoEncryptionProcessor {
968 config: AutoEncryptionConfig,
970 workspace_manager: WorkspaceKeyManager,
972 workspace_id: String,
974}
975
976impl AutoEncryptionProcessor {
977 pub fn new(workspace_id: &str, config: AutoEncryptionConfig) -> Self {
979 Self {
980 config,
981 workspace_manager: WorkspaceKeyManager::new(),
982 workspace_id: workspace_id.to_string(),
983 }
984 }
985
986 pub fn process_headers(
988 &self,
989 headers: &mut std::collections::HashMap<String, String>,
990 ) -> Result<()> {
991 if !self.config.enabled {
992 return Ok(());
993 }
994
995 let workspace_key = self.workspace_manager.get_workspace_key(&self.workspace_id)?;
996
997 for (key, value) in headers.iter_mut() {
998 if self.is_sensitive_header(key) && !self.is_already_encrypted(value) {
999 *value = workspace_key.encrypt_chacha20(value, Some(key.as_bytes()))?;
1000 }
1001 }
1002
1003 Ok(())
1004 }
1005
1006 pub fn process_json(&self, json: &mut serde_json::Value) -> Result<()> {
1008 if !self.config.enabled {
1009 return Ok(());
1010 }
1011
1012 let workspace_key = self.workspace_manager.get_workspace_key(&self.workspace_id)?;
1013 self.process_json_recursive(json, &workspace_key, Vec::new())?;
1014
1015 Ok(())
1016 }
1017
1018 pub fn process_env_vars(
1020 &self,
1021 env_vars: &mut std::collections::HashMap<String, String>,
1022 ) -> Result<()> {
1023 if !self.config.enabled {
1024 return Ok(());
1025 }
1026
1027 let workspace_key = self.workspace_manager.get_workspace_key(&self.workspace_id)?;
1028
1029 for (key, value) in env_vars.iter_mut() {
1030 if self.is_sensitive_env_var(key) && !self.is_already_encrypted(value) {
1031 *value = workspace_key.encrypt_chacha20(value, Some(key.as_bytes()))?;
1032 }
1033 }
1034
1035 Ok(())
1036 }
1037
1038 fn is_sensitive_header(&self, header_name: &str) -> bool {
1040 self.config
1041 .sensitive_headers
1042 .iter()
1043 .any(|h| h.eq_ignore_ascii_case(header_name))
1044 }
1045
1046 fn is_sensitive_env_var(&self, var_name: &str) -> bool {
1048 self.config.sensitive_env_vars.iter().any(|v| v.eq_ignore_ascii_case(var_name))
1049 }
1050
1051 fn is_sensitive_field(&self, field_path: &[String]) -> bool {
1053 let default_field = String::new();
1054 let field_name = field_path.last().unwrap_or(&default_field);
1055
1056 if self.config.sensitive_fields.iter().any(|f| f.eq_ignore_ascii_case(field_name)) {
1058 return true;
1059 }
1060
1061 let path_str = field_path.join(".");
1063 for pattern in &self.config.sensitive_patterns {
1064 if let Ok(regex) = regex::Regex::new(pattern) {
1065 if regex.is_match(&path_str) || regex.is_match(field_name) {
1066 return true;
1067 }
1068 }
1069 }
1070
1071 false
1072 }
1073
1074 fn is_already_encrypted(&self, value: &str) -> bool {
1076 value.len() > 100 && general_purpose::STANDARD.decode(value).is_ok()
1078 }
1079
1080 fn process_json_recursive(
1082 &self,
1083 json: &mut serde_json::Value,
1084 workspace_key: &EncryptionKey,
1085 current_path: Vec<String>,
1086 ) -> Result<()> {
1087 match json {
1088 serde_json::Value::Object(obj) => {
1089 for (key, value) in obj.iter_mut() {
1090 let mut new_path = current_path.clone();
1091 new_path.push(key.clone());
1092
1093 if let serde_json::Value::String(ref mut s) = value {
1094 if self.is_sensitive_field(&new_path) && !self.is_already_encrypted(s) {
1095 let path_str = new_path.join(".");
1096 let path_bytes = path_str.as_bytes();
1097 *s = workspace_key.encrypt_chacha20(s, Some(path_bytes))?;
1098 }
1099 } else {
1100 self.process_json_recursive(value, workspace_key, new_path)?;
1101 }
1102 }
1103 }
1104 serde_json::Value::Array(arr) => {
1105 for (index, item) in arr.iter_mut().enumerate() {
1106 let mut new_path = current_path.clone();
1107 new_path.push(index.to_string());
1108 self.process_json_recursive(item, workspace_key, new_path)?;
1109 }
1110 }
1111 _ => {} }
1113
1114 Ok(())
1115 }
1116}
1117
1118pub mod utils {
1120 use super::*;
1121
1122 pub async fn is_encryption_enabled_for_workspace(
1124 persistence: &WorkspacePersistence,
1125 workspace_id: &str,
1126 ) -> Result<bool> {
1127 if let Ok(workspace) = persistence.load_workspace(workspace_id).await {
1129 return Ok(workspace.config.auto_encryption.enabled);
1130 }
1131 let manager = WorkspaceKeyManager::new();
1133 Ok(manager.has_workspace_key(workspace_id))
1134 }
1135
1136 pub async fn get_auto_encryption_config(
1138 persistence: &WorkspacePersistence,
1139 workspace_id: &str,
1140 ) -> Result<AutoEncryptionConfig> {
1141 let workspace = persistence.load_workspace(workspace_id).await.map_err(|e| {
1142 EncryptionError::Generic {
1143 message: format!("Failed to load workspace: {}", e),
1144 }
1145 })?;
1146 Ok(workspace.config.auto_encryption)
1147 }
1148
1149 pub fn encrypt_for_workspace(workspace_id: &str, data: &str) -> Result<String> {
1151 let manager = WorkspaceKeyManager::new();
1152 let key = manager.get_workspace_key(workspace_id)?;
1153 key.encrypt_chacha20(data, None)
1154 }
1155
1156 pub fn decrypt_for_workspace(workspace_id: &str, encrypted_data: &str) -> Result<String> {
1158 let manager = WorkspaceKeyManager::new();
1159 let key = manager.get_workspace_key(workspace_id)?;
1160 key.decrypt_chacha20(encrypted_data, None)
1161 }
1162}
1163
1164pub fn encrypt_with_key(
1177 key_id: &str,
1178 plaintext: &str,
1179 associated_data: Option<&[u8]>,
1180) -> Result<String> {
1181 let store = get_key_store().ok_or_else(|| EncryptionError::InvalidKey {
1182 message: "Key store not initialized".to_string(),
1183 })?;
1184
1185 let key = store.get_key(key_id).ok_or_else(|| EncryptionError::InvalidKey {
1186 message: format!("Key '{}' not found", key_id),
1187 })?;
1188
1189 key.encrypt(plaintext, associated_data)
1190}
1191
1192pub fn decrypt_with_key(
1205 key_id: &str,
1206 ciphertext: &str,
1207 associated_data: Option<&[u8]>,
1208) -> Result<String> {
1209 let store = get_key_store().ok_or_else(|| EncryptionError::InvalidKey {
1210 message: "Key store not initialized".to_string(),
1211 })?;
1212
1213 let key = store.get_key(key_id).ok_or_else(|| EncryptionError::InvalidKey {
1214 message: format!("Key '{}' not found", key_id),
1215 })?;
1216
1217 key.decrypt(ciphertext, associated_data)
1218}
1219
1220mod algorithms;
1221#[allow(dead_code)]
1222mod auto_encryption;
1223mod derivation;
1224mod errors;
1225mod key_management;
1226#[allow(dead_code)]
1227mod key_rotation;
1228
1229#[cfg(test)]
1230mod tests {
1231 use super::*;
1232 use once_cell::sync::Lazy;
1233 use std::sync::Mutex;
1234
1235 static MASTER_KEY_TEST_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
1236
1237 #[test]
1238 fn test_aes_gcm_encrypt_decrypt() {
1239 let key = EncryptionKey::from_password_pbkdf2(
1240 "test_password",
1241 None,
1242 EncryptionAlgorithm::Aes256Gcm,
1243 )
1244 .unwrap();
1245
1246 let plaintext = "Hello, World!";
1247 let ciphertext = key.encrypt(plaintext, None).unwrap();
1248 let decrypted = key.decrypt(&ciphertext, None).unwrap();
1249
1250 assert_eq!(plaintext, decrypted);
1251 }
1252
1253 #[test]
1254 fn test_chacha20_encrypt_decrypt() {
1255 let key = EncryptionKey::from_password_pbkdf2(
1256 "test_password",
1257 None,
1258 EncryptionAlgorithm::ChaCha20Poly1305,
1259 )
1260 .unwrap();
1261
1262 let plaintext = "Hello, World!";
1263 let ciphertext = key.encrypt(plaintext, None).unwrap();
1264 let decrypted = key.decrypt(&ciphertext, None).unwrap();
1265
1266 assert_eq!(plaintext, decrypted);
1267 }
1268
1269 #[test]
1270 fn test_key_store() {
1271 let mut store = KeyStore::new();
1272
1273 store
1274 .derive_and_store_key(
1275 "test_key".to_string(),
1276 "test_password",
1277 EncryptionAlgorithm::Aes256Gcm,
1278 KeyDerivationMethod::Pbkdf2,
1279 )
1280 .unwrap();
1281
1282 assert!(store.get_key("test_key").is_some());
1283 assert!(store.list_keys().contains(&"test_key".to_string()));
1284
1285 store.remove_key("test_key");
1286 assert!(store.get_key("test_key").is_none());
1287 }
1288
1289 #[test]
1290 fn test_invalid_key_length() {
1291 let result = EncryptionKey::new(EncryptionAlgorithm::Aes256Gcm, vec![1, 2, 3]);
1292 assert!(matches!(result, Err(EncryptionError::InvalidKey { message: _ })));
1293 }
1294
1295 #[test]
1296 fn test_invalid_ciphertext() {
1297 let key = EncryptionKey::from_password_pbkdf2("test", None, EncryptionAlgorithm::Aes256Gcm)
1298 .unwrap();
1299 let result = key.decrypt("invalid_base64!", None);
1300 assert!(matches!(result, Err(EncryptionError::InvalidCiphertext { message: _ })));
1301 }
1302
1303 #[test]
1304 fn test_chacha20_encrypt_decrypt_12byte_nonce() {
1305 let key = EncryptionKey::from_password_pbkdf2(
1306 "test_password",
1307 None,
1308 EncryptionAlgorithm::ChaCha20Poly1305,
1309 )
1310 .unwrap();
1311
1312 let plaintext = "Hello, World! This is a test of ChaCha20-Poly1305 with 12-byte nonce.";
1313 let ciphertext = key.encrypt_chacha20(plaintext, None).unwrap();
1314 let decrypted = key.decrypt_chacha20(&ciphertext, None).unwrap();
1315
1316 assert_eq!(plaintext, decrypted);
1317 }
1318
1319 #[test]
1320 fn test_secure_function_template() {
1321 use crate::templating::expand_str;
1322
1323 let template = r#"{{secure "test message"}}"#;
1325 let result = expand_str(template);
1326
1327 assert_ne!(result, "test message");
1329 assert!(!result.is_empty());
1330
1331 assert!(general_purpose::STANDARD.decode(&result).is_ok());
1333 }
1334
1335 #[test]
1336 fn test_master_key_manager() {
1337 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1338
1339 let manager = MasterKeyManager::new();
1340
1341 #[cfg(any(target_os = "macos", target_os = "linux"))]
1343 {
1344 if let Ok(home) = std::env::var("HOME") {
1345 let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
1346 let _ = std::fs::remove_file(&key_path);
1347 }
1348 }
1349
1350 assert!(!manager.has_master_key());
1352
1353 manager.generate_master_key().unwrap();
1355 assert!(manager.has_master_key());
1356
1357 let master_key = manager.get_master_key().unwrap();
1359 assert_eq!(master_key.algorithm, EncryptionAlgorithm::ChaCha20Poly1305);
1360 }
1361
1362 #[test]
1363 fn test_workspace_key_manager() {
1364 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1365
1366 let master_manager = MasterKeyManager::new();
1368 let needs_generation =
1369 !master_manager.has_master_key() || master_manager.get_master_key().is_err();
1370 if needs_generation && master_manager.generate_master_key().is_err() {
1371 return;
1373 }
1374
1375 let workspace_manager = WorkspaceKeyManager::new();
1376 let workspace_id = &format!("test_workspace_{}", uuid::Uuid::new_v4());
1377
1378 assert!(!workspace_manager.has_workspace_key(workspace_id));
1380
1381 let encrypted_key = workspace_manager.generate_workspace_key(workspace_id).unwrap();
1383 assert!(workspace_manager.has_workspace_key(workspace_id));
1384 assert!(!encrypted_key.is_empty());
1385
1386 let workspace_key = workspace_manager.get_workspace_key(workspace_id).unwrap();
1388 assert_eq!(workspace_key.algorithm, EncryptionAlgorithm::ChaCha20Poly1305);
1389
1390 let test_data = "sensitive workspace data";
1392 let ciphertext = workspace_key.encrypt_chacha20(test_data, None).unwrap();
1393 let decrypted = workspace_key.decrypt_chacha20(&ciphertext, None).unwrap();
1394 assert_eq!(test_data, decrypted);
1395 }
1396
1397 #[test]
1398 fn test_backup_string_formatting() {
1399 let manager = WorkspaceKeyManager::new();
1400
1401 let test_key = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrst";
1403 let backup = manager.format_backup_string(test_key);
1404
1405 assert!(backup.contains('-'));
1407
1408 let parsed = manager.parse_backup_string(&backup).unwrap();
1410 assert_eq!(parsed, test_key);
1411 }
1412
1413 #[test]
1414 fn test_auto_encryption_processor() {
1415 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1416
1417 let master_manager = MasterKeyManager::new();
1419 let needs_key =
1420 !master_manager.has_master_key() || master_manager.get_master_key().is_err();
1421 if needs_key && master_manager.generate_master_key().is_err() {
1422 eprintln!("Skipping test_auto_encryption_processor: Failed to generate master key");
1423 return;
1424 }
1425
1426 let workspace_manager = WorkspaceKeyManager::new();
1427 let workspace_id = &format!("test_auto_encrypt_workspace_{}", uuid::Uuid::new_v4());
1428
1429 workspace_manager.generate_workspace_key(workspace_id).unwrap();
1430
1431 let config = AutoEncryptionConfig {
1432 enabled: true,
1433 ..AutoEncryptionConfig::default()
1434 };
1435
1436 let processor = AutoEncryptionProcessor::new(workspace_id, config);
1437
1438 let mut headers = std::collections::HashMap::new();
1440 headers.insert("Authorization".to_string(), "Bearer my-secret-token".to_string());
1441 headers.insert("Content-Type".to_string(), "application/json".to_string());
1442
1443 processor.process_headers(&mut headers).unwrap();
1444
1445 assert_ne!(headers["Authorization"], "Bearer my-secret-token");
1447 assert!(general_purpose::STANDARD.decode(&headers["Authorization"]).is_ok());
1448
1449 assert_eq!(headers["Content-Type"], "application/json");
1451 }
1452
1453 #[test]
1454 #[ignore] fn test_json_field_encryption() {
1456 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1457
1458 let master_manager = MasterKeyManager::new();
1460 if !master_manager.has_master_key() || master_manager.get_master_key().is_err() {
1461 master_manager.generate_master_key().unwrap();
1462 }
1463
1464 let workspace_manager = WorkspaceKeyManager::new();
1465 let workspace_id = &format!("test_json_workspace_{}", uuid::Uuid::new_v4());
1466
1467 workspace_manager.generate_workspace_key(workspace_id).unwrap();
1468
1469 let another_manager = WorkspaceKeyManager::new();
1471 assert!(another_manager.has_workspace_key(workspace_id));
1472
1473 let config = AutoEncryptionConfig {
1474 enabled: true,
1475 ..AutoEncryptionConfig::default()
1476 };
1477
1478 let processor = AutoEncryptionProcessor::new(workspace_id, config);
1479
1480 let mut json = serde_json::json!({
1482 "username": "testuser",
1483 "password": "secret123",
1484 "email": "test@example.com",
1485 "nested": {
1486 "token": "my-api-token",
1487 "normal_field": "normal_value"
1488 }
1489 });
1490
1491 processor.process_json(&mut json).unwrap();
1492
1493 assert_ne!(json["password"], "secret123");
1495 assert_ne!(json["nested"]["token"], "my-api-token");
1496
1497 assert_eq!(json["username"], "testuser");
1499 assert_eq!(json["email"], "test@example.com");
1500 assert_eq!(json["nested"]["normal_field"], "normal_value");
1501 }
1502
1503 #[test]
1504 fn test_env_var_encryption() {
1505 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1506
1507 let master_manager = MasterKeyManager::new();
1509 if !master_manager.has_master_key() || master_manager.get_master_key().is_err() {
1510 master_manager.generate_master_key().unwrap();
1511 }
1512
1513 let workspace_manager = WorkspaceKeyManager::new();
1514 let workspace_id = &format!("test_env_workspace_{}", uuid::Uuid::new_v4());
1515
1516 workspace_manager.generate_workspace_key(workspace_id).unwrap();
1517
1518 let config = AutoEncryptionConfig {
1519 enabled: true,
1520 ..AutoEncryptionConfig::default()
1521 };
1522
1523 let processor = AutoEncryptionProcessor::new(workspace_id, config);
1524
1525 let mut env_vars = std::collections::HashMap::new();
1527 env_vars.insert("API_KEY".to_string(), "sk-1234567890abcdef".to_string());
1528 env_vars
1529 .insert("DATABASE_URL".to_string(), "postgres://user:pass@host:5432/db".to_string());
1530 env_vars.insert("NORMAL_VAR".to_string(), "normal_value".to_string());
1531
1532 processor.process_env_vars(&mut env_vars).unwrap();
1533
1534 assert_ne!(env_vars["API_KEY"], "sk-1234567890abcdef");
1536 assert_ne!(env_vars["DATABASE_URL"], "postgres://user:pass@host:5432/db");
1537
1538 assert_eq!(env_vars["NORMAL_VAR"], "normal_value");
1540 }
1541
1542 #[test]
1543 fn test_encryption_utils() {
1544 let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
1545
1546 let master_manager = MasterKeyManager::new();
1548 let needs_key =
1549 !master_manager.has_master_key() || master_manager.get_master_key().is_err();
1550 if needs_key && master_manager.generate_master_key().is_err() {
1551 eprintln!("Skipping test_encryption_utils: Failed to generate master key");
1552 return;
1553 }
1554
1555 let workspace_manager = WorkspaceKeyManager::new();
1556 let workspace_id = &format!("test_utils_workspace_{}", uuid::Uuid::new_v4());
1557 workspace_manager.generate_workspace_key(workspace_id).unwrap();
1558
1559 assert!(workspace_manager.has_workspace_key(workspace_id));
1561
1562 let another_manager = WorkspaceKeyManager::new();
1564 assert!(another_manager.has_workspace_key(workspace_id));
1565
1566 let test_data = "test data for utils";
1567 let encrypted = utils::encrypt_for_workspace(workspace_id, test_data).unwrap();
1568 let decrypted = utils::decrypt_for_workspace(workspace_id, &encrypted).unwrap();
1569
1570 assert_eq!(test_data, decrypted);
1571 assert_ne!(encrypted, test_data);
1572 }
1573}