diff --git a/app/api-token-vault/src/vault.rs b/app/api-token-vault/src/vault.rs new file mode 100644 index 0000000..ab55c6d --- /dev/null +++ b/app/api-token-vault/src/vault.rs @@ -0,0 +1,338 @@ +use crate::token::{TokenData, TokenGenerator, generate_token_value}; +use crate::crypto::{CryptoManager, CryptoError, generate_salt, salt_to_base64, salt_from_base64}; +use chrono::{DateTime, Utc}; +use serde::{Serialize, Deserialize}; +use std::collections::HashMap; +use std::fs; +use std::path::{Path, PathBuf}; +use dirs; + +#[derive(Debug, Serialize, Deserialize)] +pub struct Vault { + pub project_name: String, + pub created_at: DateTime, + pub updated_at: DateTime, + pub salt: String, + pub tokens: HashMap, + version: u32, +} + +#[derive(Debug)] +pub enum VaultError { + NotInitialized, + InvalidPassword, + CorruptedData, + FileNotFound(String), + IoError(String), + TokenNotFound(String), + TokenAlreadyExists(String), +} + +impl std::fmt::Display for VaultError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + VaultError::NotInitialized => write!(f, "Vault not initialized. Run 'init' first."), + VaultError::InvalidPassword => write!(f, "Invalid master password"), + VaultError::CorruptedData => write!(f, "Vault file is corrupted"), + VaultError::FileNotFound(path) => write!(f, "Vault file not found: {}", path), + VaultError::IoError(msg) => write!(f, "IO error: {}", msg), + VaultError::TokenNotFound(name) => write!(f, "Token not found: {}", name), + VaultError::TokenAlreadyExists(name) => write!(f, "Token already exists: {}", name), + } + } +} + +impl std::error::Error for VaultError {} + +impl Vault { + const VERSION: u32 = 1; + + pub fn initialize(password: &str, project_name: &str) -> Result { + let salt = generate_salt(); + let salt_b64 = salt_to_base64(&salt); + + let crypto = CryptoManager::from_password(password, &salt) + .map_err(|_| VaultError::InvalidPassword)?; + + let vault = Vault { + project_name: project_name.to_string(), + created_at: Utc::now(), + updated_at: Utc::now(), + salt: salt_b64, + tokens: HashMap::new(), + version: Self::VERSION, + }; + + vault.save_to_file(password)?; + + Ok(vault) + } + + pub fn load(password: &str, project_name: &str) -> Result { + let vault_path = Self::get_vault_path(project_name)?; + + if !vault_path.exists() { + return Err(VaultError::NotInitialized); + } + + let content = fs::read_to_string(&vault_path) + .map_err(|e| VaultError::IoError(e.to_string()))?; + + let encrypted_data: EncryptedVaultData = serde_json::from_str(&content) + .map_err(|_| VaultError::CorruptedData)?; + + let salt = salt_from_base64(&encrypted_data.salt) + .map_err(|_| VaultError::CorruptedData)?; + + let crypto = CryptoManager::from_password(password, &salt) + .map_err(|_| VaultError::InvalidPassword)?; + + let json_data = crypto.decrypt_base64(&encrypted_data.encrypted_data) + .map_err(|_| VaultError::InvalidPassword)?; + + let mut vault: Vault = serde_json::from_str(&json_data) + .map_err(|_| VaultError::CorruptedData)?; + + vault.updated_at = Utc::now(); + + Ok(vault) + } + + pub fn save(&mut self, password: &str) -> Result<(), VaultError> { + self.updated_at = Utc::now(); + self.save_to_file(password) + } + + fn save_to_file(&self, password: &str) -> Result<(), VaultError> { + let salt = salt_from_base64(&self.salt) + .map_err(|_| VaultError::CorruptedData)?; + + let crypto = CryptoManager::from_password(password, &salt) + .map_err(|_| VaultError::InvalidPassword)?; + + let json_data = serde_json::to_string(self) + .map_err(|_| VaultError::CorruptedData)?; + + let encrypted_data = crypto.encrypt_base64(&json_data) + .map_err(|_| VaultError::CorruptedData)?; + + let vault_path = Self::get_vault_path(&self.project_name)?; + + if let Some(parent) = vault_path.parent() { + if !parent.exists() { + fs::create_dir_all(parent) + .map_err(|e| VaultError::IoError(e.to_string()))?; + } + } + + let encrypted_vault = EncryptedVaultData { + version: self.version, + salt: self.salt.clone(), + encrypted_data, + }; + + let content = serde_json::to_string(&encrypted_vault) + .map_err(|_| VaultError::CorruptedData)?; + + fs::write(&vault_path, content) + .map_err(|e| VaultError::IoError(e.to_string()))?; + + Ok(()) + } + + pub fn get_vault_path(project_name: &str) -> Result { + let config_dir = dirs::config_dir() + .ok_or_else(|| VaultError::IoError("Could not find config directory".to_string()))?; + + let vault_dir = config_dir.join("api-token-vault"); + let vault_file = vault_dir.join(format!("{}.json", project_name)); + + Ok(vault_file) + } + + pub fn generate_token(&mut self, name: &str, length: usize) -> Result { + if self.tokens.contains_key(name) { + return Err(VaultError::TokenAlreadyExists(name.to_string())); + } + + let value = generate_token_value(Some(length), false); + let token_data = TokenData::new(value.clone(), None); + + self.tokens.insert(name.to_string(), token_data); + + Ok(value) + } + + pub fn get_token(&self, name: &str) -> Option<&TokenData> { + self.tokens.get(name) + } + + pub fn remove_token(&mut self, name: &str) -> Result<(), VaultError> { + if !self.tokens.contains_key(name) { + return Err(VaultError::TokenNotFound(name.to_string())); + } + + self.tokens.remove(name); + Ok(()) + } + + pub fn rotate_token(&mut self, name: &str, force: bool) -> Result { + let token_data = self.tokens.get_mut(name) + .ok_or_else(|| VaultError::TokenNotFound(name.to_string()))?; + + if !force && !token_data.should_rotate() { + return Err(VaultError::TokenNotFound(name.to_string())); + } + + let new_value = generate_token_value(Some(32), false); + let rotation_days = token_data.rotation_days; + + *token_data = TokenData::new(new_value.clone(), rotation_days); + token_data.last_rotated = Some(Utc::now()); + + Ok(new_value) + } + + pub fn set_rotation(&mut self, name: &str, days: u32) -> Result<(), VaultError> { + let token_data = self.tokens.get_mut(name) + .ok_or_else(|| VaultError::TokenNotFound(name.to_string()))?; + + token_data.auto_rotate = true; + token_data.rotation_days = Some(days); + token_data.expires_at = Some(Utc::now() + chrono::Duration::days(days as i64)); + + Ok(()) + } + + pub fn check_expired_tokens(&self) -> HashMap { + self.tokens.iter() + .filter(|(_, token)| token.is_expired()) + .map(|(k, v)| (k.clone(), v)) + .collect() + } + + pub fn rotate_expired_tokens(&mut self) -> Vec { + let expired: Vec = self.tokens.iter() + .filter(|(_, token)| token.should_rotate()) + .map(|(k, _)| k.clone()) + .collect(); + + let mut rotated = Vec::new(); + + for name in &expired { + if let Ok(new_value) = self.rotate_token(name, true) { + println!("Rotated {}: {}", name, new_value); + rotated.push(name.clone()); + } + } + + rotated + } + + pub fn add_token(&mut self, name: &str, value: &str, rotation_days: Option) -> Result<(), VaultError> { + if self.tokens.contains_key(name) { + return Err(VaultError::TokenAlreadyExists(name.to_string())); + } + + let token_data = TokenData::new(value.to_string(), rotation_days); + self.tokens.insert(name.to_string(), token_data); + + Ok(()) + } +} + +#[derive(Debug, Serialize, Deserialize)] +struct EncryptedVaultData { + version: u32, + salt: String, + encrypted_data: String, +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use std::path::PathBuf; + + fn get_test_vault_path() -> PathBuf { + PathBuf::from("/tmp/test-api-token-vault") + } + + #[test] + fn test_vault_initialization() { + let password = "test_password"; + let project = "test_project"; + + let vault = Vault::initialize(password, project); + assert!(vault.is_ok()); + + let vault = vault.unwrap(); + assert_eq!(vault.project_name, project); + assert!(vault.tokens.is_empty()); + assert!(!vault.salt.is_empty()); + } + + #[test] + fn test_vault_save_and_load() { + let password = "test_password"; + let project = "test_load_project"; + + let mut vault = Vault::initialize(password, project).unwrap(); + vault.generate_token("test_token", 32).unwrap(); + + let vault_path = Vault::get_vault_path(project).unwrap(); + if vault_path.exists() { + fs::remove_file(&vault_path).unwrap(); + } + + vault.save_to_file(password).unwrap(); + + let loaded_vault = Vault::load(password, project); + assert!(loaded_vault.is_ok()); + + let loaded_vault = loaded_vault.unwrap(); + assert_eq!(loaded_vault.tokens.len(), 1); + assert!(loaded_vault.tokens.contains_key("test_token")); + + if vault_path.exists() { + fs::remove_file(&vault_path).unwrap(); + } + } + + #[test] + fn test_wrong_password() { + let password = "correct_password"; + let wrong_password = "wrong_password"; + let project = "test_wrong_password"; + + let vault = Vault::initialize(password, project).unwrap(); + let vault_path = Vault::get_vault_path(project).unwrap(); + vault.save_to_file(password).unwrap(); + + let loaded = Vault::load(wrong_password, project); + assert!(loaded.is_err()); + + if vault_path.exists() { + fs::remove_file(&vault_path).unwrap(); + } + } + + #[test] + fn test_token_rotation() { + let password = "test_password"; + let project = "test_rotation"; + + let mut vault = Vault::initialize(password, project).unwrap(); + vault.generate_token("rotate_me", 32).unwrap(); + + let original_token = vault.get_token("rotate_me").unwrap().value.clone(); + let new_value = vault.rotate_token("rotate_me", true).unwrap(); + + assert_ne!(original_token, new_value); + + let vault_path = Vault::get_vault_path(project).unwrap(); + if vault_path.exists() { + fs::remove_file(&vault_path).unwrap(); + } + } +}