pub use errors::*;
pub use key_management::{FileKeyStorage, KeyStorage, KeyStore as KeyManagementStore};
use crate::workspace_persistence::WorkspacePersistence;
use aes_gcm::{
aead::{Aead, KeyInit},
Aes256Gcm, Nonce,
};
use argon2::{
password_hash::{PasswordHasher, SaltString},
Argon2, Params,
};
use base64::{engine::general_purpose, Engine as _};
use chacha20poly1305::{ChaCha20Poly1305, Key as ChaChaKey};
use pbkdf2::pbkdf2_hmac;
use rand::{thread_rng, Rng};
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use std::fmt;
use thiserror::Error;
use tracing;
#[derive(Error, Debug)]
pub enum EncryptionError {
#[error("Encryption failure: {0}")]
Encryption(String),
#[error("Decryption failure: {0}")]
Decryption(String),
#[error("Invalid key: {0}")]
InvalidKey(String),
#[error("Invalid ciphertext: {0}")]
InvalidCiphertext(String),
#[error("Key derivation failure: {0}")]
KeyDerivation(String),
#[error("Generic encryption error: {message}")]
Generic {
message: String,
},
}
pub type Result<T> = std::result::Result<T, EncryptionError>;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EncryptionAlgorithm {
Aes256Gcm,
ChaCha20Poly1305,
}
impl fmt::Display for EncryptionAlgorithm {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
EncryptionAlgorithm::Aes256Gcm => write!(f, "aes256-gcm"),
EncryptionAlgorithm::ChaCha20Poly1305 => write!(f, "chacha20-poly1305"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KeyDerivationMethod {
Pbkdf2,
Argon2,
}
impl fmt::Display for KeyDerivationMethod {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
KeyDerivationMethod::Pbkdf2 => write!(f, "pbkdf2"),
KeyDerivationMethod::Argon2 => write!(f, "argon2"),
}
}
}
pub struct EncryptionKey {
algorithm: EncryptionAlgorithm,
key_data: Vec<u8>,
}
impl EncryptionKey {
pub fn new(algorithm: EncryptionAlgorithm, key_data: Vec<u8>) -> Result<Self> {
let expected_len = match algorithm {
EncryptionAlgorithm::Aes256Gcm => 32, EncryptionAlgorithm::ChaCha20Poly1305 => 32, };
if key_data.len() != expected_len {
return Err(EncryptionError::InvalidKey(format!(
"Key must be {} bytes for {}, got {}",
expected_len,
algorithm,
key_data.len()
)));
}
Ok(Self {
algorithm,
key_data,
})
}
pub fn from_password_pbkdf2(
password: &str,
salt: Option<&[u8]>,
algorithm: EncryptionAlgorithm,
) -> Result<Self> {
let salt = salt
.map(|s| s.to_vec())
.unwrap_or_else(|| thread_rng().random::<[u8; 32]>().to_vec());
let mut key = vec![0u8; 32];
pbkdf2_hmac::<Sha256>(password.as_bytes(), &salt, 100_000, &mut key);
Self::new(algorithm, key)
}
pub fn from_password_argon2(
password: &str,
salt: Option<&[u8]>,
algorithm: EncryptionAlgorithm,
) -> Result<Self> {
let salt_string = if let Some(salt) = salt {
SaltString::encode_b64(salt)
.map_err(|e| EncryptionError::KeyDerivation(e.to_string()))?
} else {
let mut salt_bytes = [0u8; 32];
thread_rng().fill(&mut salt_bytes);
SaltString::encode_b64(&salt_bytes)
.map_err(|e| EncryptionError::KeyDerivation(e.to_string()))?
};
let argon2 = Argon2::new(
argon2::Algorithm::Argon2id,
argon2::Version::V0x13,
Params::new(65536, 3, 1, Some(32))
.map_err(|e| EncryptionError::KeyDerivation(e.to_string()))?,
);
let hash = argon2
.hash_password(password.as_bytes(), &salt_string)
.map_err(|e| EncryptionError::KeyDerivation(e.to_string()))?;
let key_bytes = hash
.hash
.ok_or_else(|| EncryptionError::KeyDerivation("Failed to derive key hash".to_string()))?
.as_bytes()
.to_vec();
Self::new(algorithm, key_bytes)
}
pub fn encrypt(&self, plaintext: &str, associated_data: Option<&[u8]>) -> Result<String> {
match self.algorithm {
EncryptionAlgorithm::Aes256Gcm => self.encrypt_aes_gcm(plaintext, associated_data),
EncryptionAlgorithm::ChaCha20Poly1305 => {
self.encrypt_chacha20(plaintext, associated_data)
}
}
}
pub fn decrypt(&self, ciphertext: &str, associated_data: Option<&[u8]>) -> Result<String> {
match self.algorithm {
EncryptionAlgorithm::Aes256Gcm => self.decrypt_aes_gcm(ciphertext, associated_data),
EncryptionAlgorithm::ChaCha20Poly1305 => {
self.decrypt_chacha20(ciphertext, associated_data)
}
}
}
fn encrypt_aes_gcm(&self, plaintext: &str, associated_data: Option<&[u8]>) -> Result<String> {
let key_array: [u8; 32] =
self.key_data.as_slice().try_into().map_err(|_| {
EncryptionError::InvalidKey("Key length must be 32 bytes".to_string())
})?;
let cipher = Aes256Gcm::new(&key_array.into());
let nonce: [u8; 12] = thread_rng().random(); let nonce = Nonce::from(nonce);
let ciphertext = cipher
.encrypt(&nonce, plaintext.as_bytes())
.map_err(|e| EncryptionError::Encryption(e.to_string()))?;
let mut result = nonce.to_vec();
result.extend_from_slice(&ciphertext);
if let Some(aad) = associated_data {
result.extend_from_slice(aad);
}
Ok(general_purpose::STANDARD.encode(&result))
}
fn decrypt_aes_gcm(&self, ciphertext: &str, associated_data: Option<&[u8]>) -> Result<String> {
let data = general_purpose::STANDARD
.decode(ciphertext)
.map_err(|e| EncryptionError::InvalidCiphertext(e.to_string()))?;
if data.len() < 12 {
return Err(EncryptionError::InvalidCiphertext("Ciphertext too short".to_string()));
}
let nonce_array: [u8; 12] = data[0..12].try_into().map_err(|_| {
EncryptionError::InvalidCiphertext("Nonce must be 12 bytes".to_string())
})?;
let nonce = Nonce::from(nonce_array);
let ciphertext_len = if let Some(aad) = &associated_data {
let aad_len = aad.len();
data.len() - 12 - aad_len
} else {
data.len() - 12
};
let ciphertext = &data[12..12 + ciphertext_len];
let key_array: [u8; 32] =
self.key_data.as_slice().try_into().map_err(|_| {
EncryptionError::InvalidKey("Key length must be 32 bytes".to_string())
})?;
let cipher = Aes256Gcm::new(&key_array.into());
let plaintext = cipher
.decrypt(&nonce, ciphertext.as_ref())
.map_err(|e| EncryptionError::Decryption(e.to_string()))?;
String::from_utf8(plaintext)
.map_err(|e| EncryptionError::Decryption(format!("Invalid UTF-8: {}", e)))
}
pub fn encrypt_chacha20(
&self,
plaintext: &str,
_associated_data: Option<&[u8]>,
) -> Result<String> {
let key = ChaChaKey::from_slice(&self.key_data);
let cipher = ChaCha20Poly1305::new(key);
let nonce: [u8; 12] = thread_rng().random(); let nonce = chacha20poly1305::Nonce::from_slice(&nonce);
let ciphertext = cipher
.encrypt(nonce, plaintext.as_bytes())
.map_err(|e| EncryptionError::Encryption(e.to_string()))?;
let mut result = nonce.to_vec();
result.extend_from_slice(&ciphertext);
Ok(general_purpose::STANDARD.encode(&result))
}
pub fn decrypt_chacha20(
&self,
ciphertext: &str,
_associated_data: Option<&[u8]>,
) -> Result<String> {
let data = general_purpose::STANDARD
.decode(ciphertext)
.map_err(|e| EncryptionError::InvalidCiphertext(e.to_string()))?;
if data.len() < 12 {
return Err(EncryptionError::InvalidCiphertext("Ciphertext too short".to_string()));
}
let nonce = chacha20poly1305::Nonce::from_slice(&data[0..12]);
let ciphertext_data = &data[12..];
let key = ChaChaKey::from_slice(&self.key_data);
let cipher = ChaCha20Poly1305::new(key);
let plaintext = cipher
.decrypt(nonce, ciphertext_data.as_ref())
.map_err(|e| EncryptionError::Decryption(e.to_string()))?;
String::from_utf8(plaintext)
.map_err(|e| EncryptionError::Decryption(format!("Invalid UTF-8: {}", e)))
}
}
pub struct KeyStore {
keys: std::collections::HashMap<String, EncryptionKey>,
}
impl KeyStore {
pub fn new() -> Self {
Self {
keys: std::collections::HashMap::new(),
}
}
pub fn store_key(&mut self, id: String, key: EncryptionKey) {
self.keys.insert(id, key);
}
pub fn get_key(&self, id: &str) -> Option<&EncryptionKey> {
self.keys.get(id)
}
pub fn remove_key(&mut self, id: &str) -> bool {
self.keys.remove(id).is_some()
}
pub fn list_keys(&self) -> Vec<String> {
self.keys.keys().cloned().collect()
}
pub fn derive_and_store_key(
&mut self,
id: String,
password: &str,
algorithm: EncryptionAlgorithm,
method: KeyDerivationMethod,
) -> Result<()> {
let key = match method {
KeyDerivationMethod::Pbkdf2 => {
EncryptionKey::from_password_pbkdf2(password, None, algorithm)?
}
KeyDerivationMethod::Argon2 => {
EncryptionKey::from_password_argon2(password, None, algorithm)?
}
};
self.store_key(id, key);
Ok(())
}
}
impl Default for KeyStore {
fn default() -> Self {
Self::new()
}
}
static KEY_STORE: once_cell::sync::OnceCell<KeyStore> = once_cell::sync::OnceCell::new();
pub fn init_key_store() -> &'static KeyStore {
KEY_STORE.get_or_init(KeyStore::default)
}
pub fn get_key_store() -> Option<&'static KeyStore> {
KEY_STORE.get()
}
pub struct MasterKeyManager {
_service_name: String,
_account_name: String,
}
impl MasterKeyManager {
pub fn new() -> Self {
Self {
_service_name: "com.mockforge.encryption".to_string(),
_account_name: "master_key".to_string(),
}
}
pub fn generate_master_key(&self) -> Result<()> {
let master_key_bytes: [u8; 32] = rand::random();
let master_key_b64 = general_purpose::STANDARD.encode(master_key_bytes);
#[cfg(target_os = "macos")]
{
self.store_in_macos_keychain(&master_key_b64)?;
}
#[cfg(target_os = "linux")]
{
self.store_in_linux_keyring(&master_key_b64)?;
}
#[cfg(target_os = "windows")]
{
self.store_in_windows_credential_manager(&master_key_b64)?;
}
Ok(())
}
pub fn get_master_key(&self) -> Result<EncryptionKey> {
let master_key_b64 = self.retrieve_from_keychain()?;
let master_key_bytes = general_purpose::STANDARD
.decode(master_key_b64)
.map_err(|e| EncryptionError::InvalidKey(e.to_string()))?;
if master_key_bytes.len() != 32 {
return Err(EncryptionError::InvalidKey("Invalid master key length".to_string()));
}
EncryptionKey::new(EncryptionAlgorithm::ChaCha20Poly1305, master_key_bytes)
}
pub fn has_master_key(&self) -> bool {
self.retrieve_from_keychain().is_ok()
}
#[cfg(target_os = "macos")]
fn store_in_macos_keychain(&self, key: &str) -> Result<()> {
use std::os::unix::fs::PermissionsExt;
let home = std::env::var("HOME").map_err(|_| {
EncryptionError::InvalidKey("HOME environment variable not set".to_string())
})?;
let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
if let Some(parent) = key_path.parent() {
std::fs::create_dir_all(parent).map_err(|e| {
EncryptionError::InvalidKey(format!("Failed to create directory: {}", e))
})?;
}
std::fs::write(&key_path, key).map_err(|e| {
EncryptionError::InvalidKey(format!("Failed to write master key: {}", e))
})?;
let mut perms = std::fs::metadata(&key_path)
.map_err(|e| EncryptionError::InvalidKey(format!("Failed to get metadata: {}", e)))?
.permissions();
perms.set_mode(0o600);
std::fs::set_permissions(&key_path, perms).map_err(|e| {
EncryptionError::InvalidKey(format!("Failed to set permissions: {}", e))
})?;
Ok(())
}
#[cfg(target_os = "linux")]
fn store_in_linux_keyring(&self, key: &str) -> Result<()> {
use std::os::unix::fs::PermissionsExt;
let home = std::env::var("HOME").map_err(|_| {
EncryptionError::InvalidKey("HOME environment variable not set".to_string())
})?;
let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
if let Some(parent) = key_path.parent() {
std::fs::create_dir_all(parent).map_err(|e| {
EncryptionError::InvalidKey(format!("Failed to create directory: {}", e))
})?;
}
std::fs::write(&key_path, key).map_err(|e| {
EncryptionError::InvalidKey(format!("Failed to write master key: {}", e))
})?;
let mut perms = std::fs::metadata(&key_path)
.map_err(|e| EncryptionError::InvalidKey(format!("Failed to get metadata: {}", e)))?
.permissions();
perms.set_mode(0o600);
std::fs::set_permissions(&key_path, perms).map_err(|e| {
EncryptionError::InvalidKey(format!("Failed to set permissions: {}", e))
})?;
Ok(())
}
#[cfg(target_os = "windows")]
fn store_in_windows_credential_manager(&self, key: &str) -> Result<()> {
use std::ffi::OsStr;
use std::os::windows::ffi::OsStrExt;
use windows::core::PWSTR;
use windows::Win32::Security::Credentials::{
CredWriteW, CREDENTIALW, CRED_FLAGS, CRED_PERSIST_LOCAL_MACHINE, CRED_TYPE_GENERIC,
};
let target_name = "MockForge/MasterKey";
let mut target_name_wide: Vec<u16> =
OsStr::new(target_name).encode_wide().chain(std::iter::once(0)).collect();
let mut credential_blob: Vec<u16> =
OsStr::new(key).encode_wide().chain(std::iter::once(0)).collect();
let credential = CREDENTIALW {
Flags: CRED_FLAGS::default(),
Type: CRED_TYPE_GENERIC,
TargetName: PWSTR::from_raw(target_name_wide.as_mut_ptr()),
Comment: PWSTR::null(),
LastWritten: windows::Win32::Foundation::FILETIME::default(),
CredentialBlobSize: (credential_blob.len() * 2) as u32,
CredentialBlob: credential_blob.as_mut_ptr() as *mut u8,
Persist: CRED_PERSIST_LOCAL_MACHINE,
AttributeCount: 0,
Attributes: std::ptr::null_mut(),
TargetAlias: PWSTR::null(),
UserName: PWSTR::null(),
};
unsafe {
CredWriteW(&credential, 0).map_err(|e| {
EncryptionError::InvalidKey(format!("Failed to store credential: {:?}", e))
})?;
}
Ok(())
}
fn retrieve_from_keychain(&self) -> Result<String> {
#[cfg(target_os = "macos")]
{
let home = std::env::var("HOME").map_err(|_| {
EncryptionError::InvalidKey("HOME environment variable not set".to_string())
})?;
let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
let key_str = std::fs::read_to_string(&key_path).map_err(|_| {
EncryptionError::InvalidKey("Master key not found in keychain".to_string())
})?;
Ok(key_str.trim().to_string())
}
#[cfg(target_os = "linux")]
{
let home = std::env::var("HOME").map_err(|_| {
EncryptionError::InvalidKey("HOME environment variable not set".to_string())
})?;
let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
let key_str = std::fs::read_to_string(&key_path).map_err(|_| {
EncryptionError::InvalidKey("Master key not found in keychain".to_string())
})?;
Ok(key_str.trim().to_string())
}
#[cfg(target_os = "windows")]
{
self.retrieve_from_windows_credential_manager()
}
#[cfg(not(any(target_os = "macos", target_os = "linux", target_os = "windows")))]
{
let key_str = std::env::var("MOCKFORGE_MASTER_KEY").map_err(|_| {
EncryptionError::InvalidKey("Master key not found in keychain".to_string())
})?;
Ok(key_str.trim().to_string())
}
}
#[cfg(target_os = "windows")]
fn retrieve_from_windows_credential_manager(&self) -> Result<String> {
use std::ffi::OsString;
use std::os::windows::ffi::{OsStrExt, OsStringExt};
use windows::core::PCWSTR;
use windows::Win32::Security::Credentials::{
CredFree, CredReadW, CREDENTIALW, CRED_TYPE_GENERIC,
};
let target_name = "MockForge/MasterKey";
let target_name_wide: Vec<u16> = std::ffi::OsStr::new(target_name)
.encode_wide()
.chain(std::iter::once(0))
.collect();
let mut credential_ptr: *mut CREDENTIALW = std::ptr::null_mut();
unsafe {
CredReadW(
PCWSTR::from_raw(target_name_wide.as_ptr()),
CRED_TYPE_GENERIC,
None,
&mut credential_ptr,
)
.map_err(|e| {
EncryptionError::InvalidKey(format!("Failed to read credential: {:?}", e))
})?;
if credential_ptr.is_null() {
return Err(EncryptionError::InvalidKey("Credential not found".to_string()));
}
let credential = &*credential_ptr;
let blob_slice = std::slice::from_raw_parts(
credential.CredentialBlob as *const u16,
credential.CredentialBlobSize as usize / 2, );
let credential_str = OsString::from_wide(blob_slice)
.to_string_lossy()
.trim_end_matches('\0')
.to_string();
CredFree(credential_ptr as *const std::ffi::c_void);
Ok(credential_str)
}
}
}
impl Default for MasterKeyManager {
fn default() -> Self {
Self::new()
}
}
pub struct WorkspaceKeyManager {
master_key_manager: MasterKeyManager,
key_storage: std::cell::RefCell<FileKeyStorage>,
}
impl WorkspaceKeyManager {
pub fn new() -> Self {
Self {
master_key_manager: MasterKeyManager::new(),
key_storage: std::cell::RefCell::new(FileKeyStorage::new()),
}
}
pub fn with_storage_path<P: AsRef<std::path::Path>>(path: P) -> Self {
Self {
master_key_manager: MasterKeyManager::new(),
key_storage: std::cell::RefCell::new(FileKeyStorage::with_path(path)),
}
}
pub fn generate_workspace_key(&self, workspace_id: &str) -> Result<String> {
let workspace_key_bytes: [u8; 32] = rand::random();
let master_key = self.master_key_manager.get_master_key()?;
let workspace_key_b64 = master_key.encrypt_chacha20(
&general_purpose::STANDARD.encode(workspace_key_bytes),
Some(workspace_id.as_bytes()),
)?;
self.store_workspace_key(workspace_id, &workspace_key_b64)?;
Ok(workspace_key_b64)
}
pub fn get_workspace_key(&self, workspace_id: &str) -> Result<EncryptionKey> {
let encrypted_key_b64 = self.retrieve_workspace_key(workspace_id)?;
let master_key = self.master_key_manager.get_master_key()?;
let decrypted_key_b64 =
master_key.decrypt_chacha20(&encrypted_key_b64, Some(workspace_id.as_bytes()))?;
let workspace_key_bytes = general_purpose::STANDARD
.decode(decrypted_key_b64)
.map_err(|e| EncryptionError::InvalidKey(e.to_string()))?;
if workspace_key_bytes.len() != 32 {
return Err(EncryptionError::InvalidKey("Invalid workspace key length".to_string()));
}
EncryptionKey::new(EncryptionAlgorithm::ChaCha20Poly1305, workspace_key_bytes)
}
pub fn has_workspace_key(&self, workspace_id: &str) -> bool {
self.retrieve_workspace_key(workspace_id).is_ok()
}
pub fn generate_workspace_key_backup(&self, workspace_id: &str) -> Result<String> {
let encrypted_key = self.retrieve_workspace_key(workspace_id)?;
let backup_string = self.format_backup_string(&encrypted_key);
Ok(backup_string)
}
pub fn restore_workspace_key_from_backup(
&self,
workspace_id: &str,
backup_string: &str,
) -> Result<()> {
let encrypted_key = self.parse_backup_string(backup_string)?;
self.store_workspace_key(workspace_id, &encrypted_key)
}
fn store_workspace_key(&self, workspace_id: &str, encrypted_key: &str) -> Result<()> {
self.key_storage
.borrow_mut()
.store_key(&workspace_id.to_string(), encrypted_key.as_bytes())
.map_err(|e| {
EncryptionError::InvalidKey(format!("Failed to store workspace key: {:?}", e))
})
}
fn retrieve_workspace_key(&self, workspace_id: &str) -> Result<String> {
match self.key_storage.borrow().retrieve_key(&workspace_id.to_string()) {
Ok(encrypted_bytes) => String::from_utf8(encrypted_bytes).map_err(|e| {
EncryptionError::InvalidKey(format!("Invalid UTF-8 in stored key: {}", e))
}),
Err(_) => {
let old_key_file = format!("workspace_{}_key.enc", workspace_id);
match std::fs::read_to_string(&old_key_file) {
Ok(encrypted_key) => {
if let Err(e) = self
.key_storage
.borrow_mut()
.store_key(&workspace_id.to_string(), encrypted_key.as_bytes())
{
tracing::warn!(
"Failed to migrate workspace key to new storage: {:?}",
e
);
} else {
let _ = std::fs::remove_file(&old_key_file);
}
Ok(encrypted_key)
}
Err(_) => Err(EncryptionError::InvalidKey(format!(
"Workspace key not found for: {}",
workspace_id
))),
}
}
}
}
fn format_backup_string(&self, encrypted_key: &str) -> String {
let chars: Vec<char> = encrypted_key.chars().collect();
let mut result = String::new();
for (i, &ch) in chars.iter().enumerate() {
if i > 0 && i % 6 == 0 && i < chars.len() - 1 {
result.push('-');
}
result.push(ch);
}
if result.len() > 59 {
result.truncate(59);
}
result
}
fn parse_backup_string(&self, backup_string: &str) -> Result<String> {
Ok(backup_string.replace("-", ""))
}
}
impl Default for WorkspaceKeyManager {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AutoEncryptionConfig {
pub enabled: bool,
pub sensitive_headers: Vec<String>,
pub sensitive_fields: Vec<String>,
pub sensitive_env_vars: Vec<String>,
pub sensitive_patterns: Vec<String>,
}
impl Default for AutoEncryptionConfig {
fn default() -> Self {
Self {
enabled: false,
sensitive_headers: vec![
"authorization".to_string(),
"x-api-key".to_string(),
"x-auth-token".to_string(),
"cookie".to_string(),
"set-cookie".to_string(),
],
sensitive_fields: vec![
"password".to_string(),
"token".to_string(),
"secret".to_string(),
"key".to_string(),
"credentials".to_string(),
],
sensitive_env_vars: vec![
"API_KEY".to_string(),
"SECRET_KEY".to_string(),
"PASSWORD".to_string(),
"TOKEN".to_string(),
"DATABASE_URL".to_string(),
],
sensitive_patterns: vec![
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(), ],
}
}
}
pub struct AutoEncryptionProcessor {
config: AutoEncryptionConfig,
workspace_manager: WorkspaceKeyManager,
workspace_id: String,
}
impl AutoEncryptionProcessor {
pub fn new(workspace_id: &str, config: AutoEncryptionConfig) -> Self {
Self {
config,
workspace_manager: WorkspaceKeyManager::new(),
workspace_id: workspace_id.to_string(),
}
}
pub fn process_headers(
&self,
headers: &mut std::collections::HashMap<String, String>,
) -> Result<()> {
if !self.config.enabled {
return Ok(());
}
let workspace_key = self.workspace_manager.get_workspace_key(&self.workspace_id)?;
for (key, value) in headers.iter_mut() {
if self.is_sensitive_header(key) && !self.is_already_encrypted(value) {
*value = workspace_key.encrypt_chacha20(value, Some(key.as_bytes()))?;
}
}
Ok(())
}
pub fn process_json(&self, json: &mut serde_json::Value) -> Result<()> {
if !self.config.enabled {
return Ok(());
}
let workspace_key = self.workspace_manager.get_workspace_key(&self.workspace_id)?;
self.process_json_recursive(json, &workspace_key, Vec::new())?;
Ok(())
}
pub fn process_env_vars(
&self,
env_vars: &mut std::collections::HashMap<String, String>,
) -> Result<()> {
if !self.config.enabled {
return Ok(());
}
let workspace_key = self.workspace_manager.get_workspace_key(&self.workspace_id)?;
for (key, value) in env_vars.iter_mut() {
if self.is_sensitive_env_var(key) && !self.is_already_encrypted(value) {
*value = workspace_key.encrypt_chacha20(value, Some(key.as_bytes()))?;
}
}
Ok(())
}
fn is_sensitive_header(&self, header_name: &str) -> bool {
self.config
.sensitive_headers
.iter()
.any(|h| h.eq_ignore_ascii_case(header_name))
}
fn is_sensitive_env_var(&self, var_name: &str) -> bool {
self.config.sensitive_env_vars.iter().any(|v| v.eq_ignore_ascii_case(var_name))
}
fn is_sensitive_field(&self, field_path: &[String]) -> bool {
let default_field = String::new();
let field_name = field_path.last().unwrap_or(&default_field);
if self.config.sensitive_fields.iter().any(|f| f.eq_ignore_ascii_case(field_name)) {
return true;
}
let path_str = field_path.join(".");
for pattern in &self.config.sensitive_patterns {
if let Ok(regex) = regex::Regex::new(pattern) {
if regex.is_match(&path_str) || regex.is_match(field_name) {
return true;
}
}
}
false
}
fn is_already_encrypted(&self, value: &str) -> bool {
value.len() > 100 && general_purpose::STANDARD.decode(value).is_ok()
}
fn process_json_recursive(
&self,
json: &mut serde_json::Value,
workspace_key: &EncryptionKey,
current_path: Vec<String>,
) -> Result<()> {
match json {
serde_json::Value::Object(obj) => {
for (key, value) in obj.iter_mut() {
let mut new_path = current_path.clone();
new_path.push(key.clone());
if let serde_json::Value::String(ref mut s) = value {
if self.is_sensitive_field(&new_path) && !self.is_already_encrypted(s) {
let path_str = new_path.join(".");
let path_bytes = path_str.as_bytes();
*s = workspace_key.encrypt_chacha20(s, Some(path_bytes))?;
}
} else {
self.process_json_recursive(value, workspace_key, new_path)?;
}
}
}
serde_json::Value::Array(arr) => {
for (index, item) in arr.iter_mut().enumerate() {
let mut new_path = current_path.clone();
new_path.push(index.to_string());
self.process_json_recursive(item, workspace_key, new_path)?;
}
}
_ => {} }
Ok(())
}
}
pub mod utils {
use super::*;
pub async fn is_encryption_enabled_for_workspace(
persistence: &WorkspacePersistence,
workspace_id: &str,
) -> Result<bool> {
if let Ok(workspace) = persistence.load_workspace(workspace_id).await {
return Ok(workspace.config.auto_encryption.enabled);
}
let manager = WorkspaceKeyManager::new();
Ok(manager.has_workspace_key(workspace_id))
}
pub async fn get_auto_encryption_config(
persistence: &WorkspacePersistence,
workspace_id: &str,
) -> Result<AutoEncryptionConfig> {
let workspace = persistence.load_workspace(workspace_id).await.map_err(|e| {
EncryptionError::Generic {
message: format!("Failed to load workspace: {}", e),
}
})?;
Ok(workspace.config.auto_encryption)
}
pub fn encrypt_for_workspace(workspace_id: &str, data: &str) -> Result<String> {
let manager = WorkspaceKeyManager::new();
let key = manager.get_workspace_key(workspace_id)?;
key.encrypt_chacha20(data, None)
}
pub fn decrypt_for_workspace(workspace_id: &str, encrypted_data: &str) -> Result<String> {
let manager = WorkspaceKeyManager::new();
let key = manager.get_workspace_key(workspace_id)?;
key.decrypt_chacha20(encrypted_data, None)
}
}
pub fn encrypt_with_key(
key_id: &str,
plaintext: &str,
associated_data: Option<&[u8]>,
) -> Result<String> {
let store = get_key_store()
.ok_or_else(|| EncryptionError::InvalidKey("Key store not initialized".to_string()))?;
let key = store
.get_key(key_id)
.ok_or_else(|| EncryptionError::InvalidKey(format!("Key '{}' not found", key_id)))?;
key.encrypt(plaintext, associated_data)
}
pub fn decrypt_with_key(
key_id: &str,
ciphertext: &str,
associated_data: Option<&[u8]>,
) -> Result<String> {
let store = get_key_store()
.ok_or_else(|| EncryptionError::InvalidKey("Key store not initialized".to_string()))?;
let key = store
.get_key(key_id)
.ok_or_else(|| EncryptionError::InvalidKey(format!("Key '{}' not found", key_id)))?;
key.decrypt(ciphertext, associated_data)
}
mod algorithms;
#[allow(dead_code)]
mod auto_encryption;
mod derivation;
mod errors;
mod key_management;
#[allow(dead_code)]
mod key_rotation;
#[cfg(test)]
mod tests {
use super::*;
use once_cell::sync::Lazy;
use std::sync::Mutex;
static MASTER_KEY_TEST_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
#[test]
fn test_aes_gcm_encrypt_decrypt() {
let key = EncryptionKey::from_password_pbkdf2(
"test_password",
None,
EncryptionAlgorithm::Aes256Gcm,
)
.unwrap();
let plaintext = "Hello, World!";
let ciphertext = key.encrypt(plaintext, None).unwrap();
let decrypted = key.decrypt(&ciphertext, None).unwrap();
assert_eq!(plaintext, decrypted);
}
#[test]
fn test_chacha20_encrypt_decrypt() {
let key = EncryptionKey::from_password_pbkdf2(
"test_password",
None,
EncryptionAlgorithm::ChaCha20Poly1305,
)
.unwrap();
let plaintext = "Hello, World!";
let ciphertext = key.encrypt(plaintext, None).unwrap();
let decrypted = key.decrypt(&ciphertext, None).unwrap();
assert_eq!(plaintext, decrypted);
}
#[test]
fn test_key_store() {
let mut store = KeyStore::new();
store
.derive_and_store_key(
"test_key".to_string(),
"test_password",
EncryptionAlgorithm::Aes256Gcm,
KeyDerivationMethod::Pbkdf2,
)
.unwrap();
assert!(store.get_key("test_key").is_some());
assert!(store.list_keys().contains(&"test_key".to_string()));
store.remove_key("test_key");
assert!(store.get_key("test_key").is_none());
}
#[test]
fn test_invalid_key_length() {
let result = EncryptionKey::new(EncryptionAlgorithm::Aes256Gcm, vec![1, 2, 3]);
assert!(matches!(result, Err(EncryptionError::InvalidKey(_))));
}
#[test]
fn test_invalid_ciphertext() {
let key = EncryptionKey::from_password_pbkdf2("test", None, EncryptionAlgorithm::Aes256Gcm)
.unwrap();
let result = key.decrypt("invalid_base64!", None);
assert!(matches!(result, Err(EncryptionError::InvalidCiphertext(_))));
}
#[test]
fn test_chacha20_encrypt_decrypt_12byte_nonce() {
let key = EncryptionKey::from_password_pbkdf2(
"test_password",
None,
EncryptionAlgorithm::ChaCha20Poly1305,
)
.unwrap();
let plaintext = "Hello, World! This is a test of ChaCha20-Poly1305 with 12-byte nonce.";
let ciphertext = key.encrypt_chacha20(plaintext, None).unwrap();
let decrypted = key.decrypt_chacha20(&ciphertext, None).unwrap();
assert_eq!(plaintext, decrypted);
}
#[test]
fn test_secure_function_template() {
use crate::templating::expand_str;
let template = r#"{{secure "test message"}}"#;
let result = expand_str(template);
assert_ne!(result, "test message");
assert!(!result.is_empty());
assert!(general_purpose::STANDARD.decode(&result).is_ok());
}
#[test]
fn test_master_key_manager() {
let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
let manager = MasterKeyManager::new();
#[cfg(any(target_os = "macos", target_os = "linux"))]
{
if let Ok(home) = std::env::var("HOME") {
let key_path = std::path::Path::new(&home).join(".mockforge").join("master_key");
let _ = std::fs::remove_file(&key_path);
}
}
assert!(!manager.has_master_key());
manager.generate_master_key().unwrap();
assert!(manager.has_master_key());
let master_key = manager.get_master_key().unwrap();
assert_eq!(master_key.algorithm, EncryptionAlgorithm::ChaCha20Poly1305);
}
#[test]
fn test_workspace_key_manager() {
let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
let master_manager = MasterKeyManager::new();
let needs_generation =
!master_manager.has_master_key() || master_manager.get_master_key().is_err();
if needs_generation && master_manager.generate_master_key().is_err() {
return;
}
let workspace_manager = WorkspaceKeyManager::new();
let workspace_id = &format!("test_workspace_{}", uuid::Uuid::new_v4());
assert!(!workspace_manager.has_workspace_key(workspace_id));
let encrypted_key = workspace_manager.generate_workspace_key(workspace_id).unwrap();
assert!(workspace_manager.has_workspace_key(workspace_id));
assert!(!encrypted_key.is_empty());
let workspace_key = workspace_manager.get_workspace_key(workspace_id).unwrap();
assert_eq!(workspace_key.algorithm, EncryptionAlgorithm::ChaCha20Poly1305);
let test_data = "sensitive workspace data";
let ciphertext = workspace_key.encrypt_chacha20(test_data, None).unwrap();
let decrypted = workspace_key.decrypt_chacha20(&ciphertext, None).unwrap();
assert_eq!(test_data, decrypted);
}
#[test]
fn test_backup_string_formatting() {
let manager = WorkspaceKeyManager::new();
let test_key = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrst";
let backup = manager.format_backup_string(test_key);
assert!(backup.contains('-'));
let parsed = manager.parse_backup_string(&backup).unwrap();
assert_eq!(parsed, test_key);
}
#[test]
fn test_auto_encryption_processor() {
let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
let master_manager = MasterKeyManager::new();
let needs_key =
!master_manager.has_master_key() || master_manager.get_master_key().is_err();
if needs_key && master_manager.generate_master_key().is_err() {
eprintln!("Skipping test_auto_encryption_processor: Failed to generate master key");
return;
}
let workspace_manager = WorkspaceKeyManager::new();
let workspace_id = &format!("test_auto_encrypt_workspace_{}", uuid::Uuid::new_v4());
workspace_manager.generate_workspace_key(workspace_id).unwrap();
let config = AutoEncryptionConfig {
enabled: true,
..AutoEncryptionConfig::default()
};
let processor = AutoEncryptionProcessor::new(workspace_id, config);
let mut headers = std::collections::HashMap::new();
headers.insert("Authorization".to_string(), "Bearer my-secret-token".to_string());
headers.insert("Content-Type".to_string(), "application/json".to_string());
processor.process_headers(&mut headers).unwrap();
assert_ne!(headers["Authorization"], "Bearer my-secret-token");
assert!(general_purpose::STANDARD.decode(&headers["Authorization"]).is_ok());
assert_eq!(headers["Content-Type"], "application/json");
}
#[test]
#[ignore] fn test_json_field_encryption() {
let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
let master_manager = MasterKeyManager::new();
if !master_manager.has_master_key() || master_manager.get_master_key().is_err() {
master_manager.generate_master_key().unwrap();
}
let workspace_manager = WorkspaceKeyManager::new();
let workspace_id = &format!("test_json_workspace_{}", uuid::Uuid::new_v4());
workspace_manager.generate_workspace_key(workspace_id).unwrap();
let another_manager = WorkspaceKeyManager::new();
assert!(another_manager.has_workspace_key(workspace_id));
let config = AutoEncryptionConfig {
enabled: true,
..AutoEncryptionConfig::default()
};
let processor = AutoEncryptionProcessor::new(workspace_id, config);
let mut json = serde_json::json!({
"username": "testuser",
"password": "secret123",
"email": "test@example.com",
"nested": {
"token": "my-api-token",
"normal_field": "normal_value"
}
});
processor.process_json(&mut json).unwrap();
assert_ne!(json["password"], "secret123");
assert_ne!(json["nested"]["token"], "my-api-token");
assert_eq!(json["username"], "testuser");
assert_eq!(json["email"], "test@example.com");
assert_eq!(json["nested"]["normal_field"], "normal_value");
}
#[test]
fn test_env_var_encryption() {
let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
let master_manager = MasterKeyManager::new();
if !master_manager.has_master_key() || master_manager.get_master_key().is_err() {
master_manager.generate_master_key().unwrap();
}
let workspace_manager = WorkspaceKeyManager::new();
let workspace_id = &format!("test_env_workspace_{}", uuid::Uuid::new_v4());
workspace_manager.generate_workspace_key(workspace_id).unwrap();
let config = AutoEncryptionConfig {
enabled: true,
..AutoEncryptionConfig::default()
};
let processor = AutoEncryptionProcessor::new(workspace_id, config);
let mut env_vars = std::collections::HashMap::new();
env_vars.insert("API_KEY".to_string(), "sk-1234567890abcdef".to_string());
env_vars
.insert("DATABASE_URL".to_string(), "postgres://user:pass@host:5432/db".to_string());
env_vars.insert("NORMAL_VAR".to_string(), "normal_value".to_string());
processor.process_env_vars(&mut env_vars).unwrap();
assert_ne!(env_vars["API_KEY"], "sk-1234567890abcdef");
assert_ne!(env_vars["DATABASE_URL"], "postgres://user:pass@host:5432/db");
assert_eq!(env_vars["NORMAL_VAR"], "normal_value");
}
#[test]
fn test_encryption_utils() {
let _guard = MASTER_KEY_TEST_LOCK.lock().unwrap();
let master_manager = MasterKeyManager::new();
let needs_key =
!master_manager.has_master_key() || master_manager.get_master_key().is_err();
if needs_key && master_manager.generate_master_key().is_err() {
eprintln!("Skipping test_encryption_utils: Failed to generate master key");
return;
}
let workspace_manager = WorkspaceKeyManager::new();
let workspace_id = &format!("test_utils_workspace_{}", uuid::Uuid::new_v4());
workspace_manager.generate_workspace_key(workspace_id).unwrap();
assert!(workspace_manager.has_workspace_key(workspace_id));
let another_manager = WorkspaceKeyManager::new();
assert!(another_manager.has_workspace_key(workspace_id));
let test_data = "test data for utils";
let encrypted = utils::encrypt_for_workspace(workspace_id, test_data).unwrap();
let decrypted = utils::decrypt_for_workspace(workspace_id, &encrypted).unwrap();
assert_eq!(test_data, decrypted);
assert_ne!(encrypted, test_data);
}
}