use std::collections::BTreeMap;
use std::path::PathBuf;
use color_eyre::{eyre::WrapErr, Result};
use serde::{Deserialize, Serialize};
use crate::config::AppConfig;
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CredentialEntry {
api_key: String,
}
#[cfg(unix)]
fn restrict_to_owner(path: &std::path::Path) -> Result<()> {
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))?;
Ok(())
}
#[cfg(not(unix))]
fn restrict_to_owner(_path: &std::path::Path) -> Result<()> {
Ok(())
}
pub struct CredentialStore {
path: PathBuf,
entries: BTreeMap<String, CredentialEntry>,
}
impl CredentialStore {
pub fn load(path: PathBuf) -> Result<Self> {
let entries = if path.exists() {
let content = std::fs::read_to_string(&path)
.wrap_err_with(|| format!("Reading credentials at {}", path.display()))?;
toml::from_str(&content)
.wrap_err_with(|| format!("Invalid credentials at {}", path.display()))?
} else {
BTreeMap::new()
};
Ok(Self { path, entries })
}
pub fn load_default() -> Result<Self> {
Self::load(AppConfig::config_dir()?.join("credentials.toml"))
}
pub fn get(&self, provider: &str) -> Option<&str> {
self.entries.get(provider).map(|e| e.api_key.as_str())
}
pub fn set(&mut self, provider: &str, key: String) -> Result<()> {
self.entries
.insert(provider.to_string(), CredentialEntry { api_key: key });
self.persist()
}
pub fn remove(&mut self, provider: &str) -> Result<bool> {
let existed = self.entries.remove(provider).is_some();
if existed {
self.persist()?;
}
Ok(existed)
}
pub fn providers_with_keys(&self) -> impl Iterator<Item = &str> {
self.entries.keys().map(|s| s.as_str())
}
fn persist(&self) -> Result<()> {
if let Some(parent) = self.path.parent() {
std::fs::create_dir_all(parent)?;
}
let content = toml::to_string_pretty(&self.entries)?;
let tmp = self.path.with_extension("toml.tmp");
std::fs::write(&tmp, &content)?;
restrict_to_owner(&tmp)?;
std::fs::rename(&tmp, &self.path)?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn store_at(dir: &std::path::Path) -> CredentialStore {
CredentialStore::load(dir.join("credentials.toml")).unwrap()
}
#[test]
fn set_then_get_round_trips() {
let temp = tempfile::tempdir().unwrap();
let mut store = store_at(temp.path());
store.set("anthropic", "sk-ant-123".to_string()).unwrap();
assert_eq!(store.get("anthropic"), Some("sk-ant-123"));
}
#[test]
fn a_missing_file_behaves_as_an_empty_store() {
let temp = tempfile::tempdir().unwrap();
let store = store_at(temp.path());
assert_eq!(store.get("anthropic"), None);
assert_eq!(store.providers_with_keys().count(), 0);
}
#[test]
#[cfg(unix)]
fn a_write_leaves_the_file_owner_only() {
use std::os::unix::fs::PermissionsExt;
let temp = tempfile::tempdir().unwrap();
let mut store = store_at(temp.path());
store.set("groq", "gsk-1".to_string()).unwrap();
let path = temp.path().join("credentials.toml");
let mode = std::fs::metadata(&path).unwrap().permissions().mode();
assert_eq!(mode & 0o777, 0o600);
}
#[test]
fn removing_an_unknown_provider_reports_no_entry() {
let temp = tempfile::tempdir().unwrap();
let mut store = store_at(temp.path());
assert!(!store.remove("openai").unwrap());
}
#[test]
fn removing_a_known_provider_removes_it_and_says_so() {
let temp = tempfile::tempdir().unwrap();
let mut store = store_at(temp.path());
store.set("openai", "sk-1".to_string()).unwrap();
assert!(store.remove("openai").unwrap());
assert_eq!(store.get("openai"), None);
}
#[test]
fn setting_an_existing_provider_again_overwrites_rather_than_duplicating() {
let temp = tempfile::tempdir().unwrap();
let mut store = store_at(temp.path());
store.set("xai", "old-key".to_string()).unwrap();
store.set("xai", "new-key".to_string()).unwrap();
assert_eq!(store.get("xai"), Some("new-key"));
assert_eq!(store.providers_with_keys().count(), 1);
}
#[test]
fn a_stored_key_survives_a_reload_from_disk() {
let temp = tempfile::tempdir().unwrap();
let path = temp.path().join("credentials.toml");
let mut store = CredentialStore::load(path.clone()).unwrap();
store.set("deepseek", "dsk-1".to_string()).unwrap();
let reloaded = CredentialStore::load(path).unwrap();
assert_eq!(reloaded.get("deepseek"), Some("dsk-1"));
}
}