use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
#[derive(Debug, Clone)]
pub struct SecureCredential {
provider: String,
encrypted_value: Vec<u8>,
created_at: chrono::DateTime<chrono::Utc>,
last_accessed: Arc<RwLock<chrono::DateTime<chrono::Utc>>>,
}
impl SecureCredential {
pub fn decrypt(&self, encryptor: &super::encryption::Encryptor) -> Result<String> {
let decrypted = encryptor.decrypt(&self.encrypted_value)?;
if let Ok(mut last_accessed) = self.last_accessed.write() {
*last_accessed = chrono::Utc::now();
}
Ok(String::from_utf8(decrypted)?)
}
pub fn metadata(&self) -> CredentialMetadata {
let last_accessed = self.last_accessed.read()
.map(|t| *t)
.unwrap_or(self.created_at);
CredentialMetadata {
provider: self.provider.clone(),
created_at: self.created_at,
last_accessed,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CredentialMetadata {
pub provider: String,
pub created_at: chrono::DateTime<chrono::Utc>,
pub last_accessed: chrono::DateTime<chrono::Utc>,
}
pub struct CredentialManager {
credentials: Arc<RwLock<HashMap<String, SecureCredential>>>,
encryptor: Option<super::encryption::Encryptor>,
encrypt_enabled: bool,
}
impl CredentialManager {
pub fn new(encrypt_enabled: bool) -> Result<Self> {
let encryptor = if encrypt_enabled {
Some(super::encryption::Encryptor::new()?)
} else {
None
};
Ok(Self {
credentials: Arc::new(RwLock::new(HashMap::new())),
encryptor,
encrypt_enabled,
})
}
pub fn store(&mut self, provider: &str, api_key: &str) -> Result<()> {
if provider.is_empty() {
return Err(anyhow!("Provider name cannot be empty"));
}
if api_key.is_empty() {
return Err(anyhow!("API key cannot be empty"));
}
let encrypted_value = if self.encrypt_enabled {
self.encryptor
.as_ref()
.ok_or_else(|| anyhow!("Encryptor not initialized"))?
.encrypt(api_key.as_bytes())?
} else {
api_key.as_bytes().to_vec()
};
let credential = SecureCredential {
provider: provider.to_string(),
encrypted_value,
created_at: chrono::Utc::now(),
last_accessed: Arc::new(RwLock::new(chrono::Utc::now())),
};
let mut credentials = self.credentials.write()
.map_err(|e| anyhow!("Failed to acquire write lock: {}", e))?;
credentials.insert(provider.to_string(), credential);
Ok(())
}
pub fn retrieve(&self, provider: &str) -> Result<Option<SecureCredential>> {
let credentials = self.credentials.read()
.map_err(|e| anyhow!("Failed to acquire read lock: {}", e))?;
Ok(credentials.get(provider).cloned())
}
pub fn get_api_key(&self, provider: &str) -> Result<Option<String>> {
let credential = self.retrieve(provider)?;
match credential {
Some(cred) => {
if self.encrypt_enabled {
let encryptor = self.encryptor
.as_ref()
.ok_or_else(|| anyhow!("Encryptor not initialized"))?;
Ok(Some(cred.decrypt(encryptor)?))
} else {
Ok(Some(String::from_utf8(cred.encrypted_value)?))
}
}
None => Ok(None),
}
}
pub fn remove(&mut self, provider: &str) -> Result<bool> {
let mut credentials = self.credentials.write()
.map_err(|e| anyhow!("Failed to acquire write lock: {}", e))?;
Ok(credentials.remove(provider).is_some())
}
pub fn list_metadata(&self) -> Result<Vec<CredentialMetadata>> {
let credentials = self.credentials.read()
.map_err(|e| anyhow!("Failed to acquire read lock: {}", e))?;
Ok(credentials.values()
.map(|cred| cred.metadata())
.collect())
}
pub fn rotate_encryption_key(&mut self) -> Result<()> {
if !self.encrypt_enabled {
return Ok(());
}
let old_encryptor = self.encryptor
.as_ref()
.ok_or_else(|| anyhow!("Encryptor not initialized"))?;
let new_encryptor = super::encryption::Encryptor::new()?;
let mut credentials = self.credentials.write()
.map_err(|e| anyhow!("Failed to acquire write lock: {}", e))?;
for credential in credentials.values_mut() {
let decrypted = old_encryptor.decrypt(&credential.encrypted_value)?;
credential.encrypted_value = new_encryptor.encrypt(&decrypted)?;
}
self.encryptor = Some(new_encryptor);
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_credential_storage_and_retrieval() {
let mut manager = CredentialManager::new(false).expect("should succeed");
manager.store("openai", "sk-test-key-123").expect("should succeed");
let api_key = manager.get_api_key("openai").expect("should succeed");
assert_eq!(api_key, Some("sk-test-key-123".to_string()));
}
#[test]
fn test_encrypted_credential_storage() {
let mut manager = CredentialManager::new(true).expect("should succeed");
manager.store("anthropic", "sk-ant-key-456").expect("should succeed");
let api_key = manager.get_api_key("anthropic").expect("should succeed");
assert_eq!(api_key, Some("sk-ant-key-456".to_string()));
}
#[test]
fn test_credential_removal() {
let mut manager = CredentialManager::new(false).expect("should succeed");
manager.store("test_provider", "test_key").expect("should succeed");
assert!(manager.get_api_key("test_provider").expect("should succeed").is_some());
manager.remove("test_provider").expect("should succeed");
assert!(manager.get_api_key("test_provider").expect("should succeed").is_none());
}
#[test]
fn test_list_metadata() {
let mut manager = CredentialManager::new(false).expect("should succeed");
manager.store("provider1", "key1").expect("should succeed");
manager.store("provider2", "key2").expect("should succeed");
let metadata = manager.list_metadata().expect("should succeed");
assert_eq!(metadata.len(), 2);
}
#[test]
fn test_invalid_input() {
let mut manager = CredentialManager::new(false).expect("should succeed");
assert!(manager.store("", "key").is_err());
assert!(manager.store("provider", "").is_err());
}
}