Initial upload: api-token-vault Rust CLI tool with encrypted vault storage
This commit is contained in:
338
app/api-token-vault/src/vault.rs
Normal file
338
app/api-token-vault/src/vault.rs
Normal file
@@ -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<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
pub salt: String,
|
||||
pub tokens: HashMap<String, TokenData>,
|
||||
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<Self, VaultError> {
|
||||
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<Self, VaultError> {
|
||||
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<PathBuf, VaultError> {
|
||||
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<String, VaultError> {
|
||||
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<String, VaultError> {
|
||||
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<String, &TokenData> {
|
||||
self.tokens.iter()
|
||||
.filter(|(_, token)| token.is_expired())
|
||||
.map(|(k, v)| (k.clone(), v))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn rotate_expired_tokens(&mut self) -> Vec<String> {
|
||||
let expired: Vec<String> = 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<u32>) -> 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();
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user