codetether_agent/secrets/
mod.rs1use anyhow::{Context, Result};
7use std::collections::HashMap;
8use std::sync::Arc;
9use tokio::sync::RwLock;
10use vaultrs::client::{VaultClient, VaultClientSettingsBuilder};
11use vaultrs::kv2;
12
13#[allow(dead_code)]
15const DEFAULT_SECRETS_PATH: &str = "secret/data/codetether/providers";
16
17#[derive(Clone)]
19pub struct SecretsManager {
20 client: Option<Arc<VaultClient>>,
21 pub cache: Arc<RwLock<HashMap<String, String>>>,
23 mount: String,
24 path: String,
25}
26
27impl Default for SecretsManager {
28 fn default() -> Self {
29 Self {
30 client: None,
31 cache: Arc::new(RwLock::new(HashMap::new())),
32 mount: "secret".to_string(),
33 path: "codetether/providers".to_string(),
34 }
35 }
36}
37
38impl SecretsManager {
39 pub async fn new(config: &VaultConfig) -> Result<Self> {
41 let settings = VaultClientSettingsBuilder::default()
42 .address(&config.address)
43 .token(&config.token)
44 .build()
45 .context("Failed to build Vault client settings")?;
46
47 let client = VaultClient::new(settings).context("Failed to create Vault client")?;
48
49 Ok(Self {
50 client: Some(Arc::new(client)),
51 cache: Arc::new(RwLock::new(HashMap::new())),
52 mount: config.mount.clone().unwrap_or_else(|| "secret".to_string()),
53 path: config
54 .path
55 .clone()
56 .unwrap_or_else(|| "codetether/providers".to_string()),
57 })
58 }
59
60 pub async fn from_k8s_auth(
67 address: &str,
68 role: &str,
69 mount: &str,
70 kv_mount: Option<&str>,
71 kv_path: Option<&str>,
72 ) -> Result<Self> {
73 let jwt_path = std::env::var("VAULT_K8S_SA_JWT_PATH")
74 .unwrap_or_else(|_| "/var/run/secrets/kubernetes.io/serviceaccount/token".to_string());
75
76 let jwt = tokio::fs::read_to_string(&jwt_path)
77 .await
78 .with_context(|| {
79 format!(
80 "Failed to read Kubernetes service account token from {}",
81 jwt_path
82 )
83 })?;
84 let jwt = jwt.trim().to_string();
85
86 let bootstrap_settings = VaultClientSettingsBuilder::default()
89 .address(address)
90 .token("")
91 .build()
92 .context("Failed to build bootstrap Vault client settings")?;
93 let bootstrap_client = VaultClient::new(bootstrap_settings)
94 .context("Failed to create bootstrap Vault client")?;
95
96 let auth_info = vaultrs::auth::kubernetes::login(&bootstrap_client, mount, role, &jwt)
97 .await
98 .context("Vault Kubernetes auth login failed")?;
99
100 let settings = VaultClientSettingsBuilder::default()
101 .address(address)
102 .token(&auth_info.client_token)
103 .build()
104 .context("Failed to build authenticated Vault client settings")?;
105 let client =
106 VaultClient::new(settings).context("Failed to create authenticated Vault client")?;
107
108 Ok(Self {
109 client: Some(Arc::new(client)),
110 cache: Arc::new(RwLock::new(HashMap::new())),
111 mount: kv_mount.unwrap_or("secret").to_string(),
112 path: kv_path.unwrap_or("codetether/providers").to_string(),
113 })
114 }
115
116 pub async fn from_env() -> Result<Self> {
123 let address = std::env::var("VAULT_ADDR").context("VAULT_ADDR not set")?;
124 let kv_mount = std::env::var("VAULT_MOUNT").ok();
125 let kv_path = std::env::var("VAULT_SECRETS_PATH").ok();
126
127 if let Ok(role) = std::env::var("VAULT_ROLE") {
132 let role = role.trim().to_string();
133 if !role.is_empty() {
134 let k8s_mount =
135 std::env::var("VAULT_AUTH_MOUNT").unwrap_or_else(|_| "kubernetes".to_string());
136
137 match Self::from_k8s_auth(
138 &address,
139 &role,
140 &k8s_mount,
141 kv_mount.as_deref(),
142 kv_path.as_deref(),
143 )
144 .await
145 {
146 Ok(manager) => {
147 tracing::info!(
148 role = %role,
149 mount = %k8s_mount,
150 "Authenticated to Vault via Kubernetes service account"
151 );
152 return Ok(manager);
153 }
154 Err(e) => {
155 tracing::warn!(
156 error = %e,
157 "Vault Kubernetes auth failed; falling back to VAULT_TOKEN"
158 );
159 }
160 }
161 }
162 }
163
164 let token = std::env::var("VAULT_TOKEN").context("VAULT_TOKEN not set")?;
165 let config = VaultConfig {
166 address,
167 token,
168 mount: kv_mount,
169 path: kv_path,
170 };
171
172 Self::new(&config).await
173 }
174
175 pub fn is_connected(&self) -> bool {
177 self.client.is_some()
178 }
179
180 pub async fn get_api_key(&self, provider_id: &str) -> Result<Option<String>> {
182 {
184 let cache = self.cache.read().await;
185 if let Some(key) = cache.get(provider_id) {
186 return Ok(Some(key.clone()));
187 }
188 }
189
190 let client = match &self.client {
192 Some(c) => c,
193 None => return Ok(None),
194 };
195
196 let secret_path = format!("{}/{}", self.path, provider_id);
197
198 match kv2::read::<ProviderSecrets>(client.as_ref(), &self.mount, &secret_path).await {
199 Ok(secret) => {
200 if let Some(ref api_key) = secret.api_key {
202 let mut cache = self.cache.write().await;
203 cache.insert(provider_id.to_string(), api_key.clone());
204 }
205 Ok(secret.api_key)
206 }
207 Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(None),
208 Err(e) => {
209 tracing::warn!("Failed to fetch secret for {}: {}", provider_id, e);
210 Ok(None)
211 }
212 }
213 }
214
215 pub async fn get_provider_secrets(&self, provider_id: &str) -> Result<Option<ProviderSecrets>> {
217 let client = match &self.client {
218 Some(c) => c,
219 None => return Ok(None),
220 };
221
222 let secret_path = format!("{}/{}", self.path, provider_id);
223
224 match kv2::read::<ProviderSecrets>(client.as_ref(), &self.mount, &secret_path).await {
225 Ok(secret) => Ok(Some(secret)),
226 Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(None),
227 Err(e) => {
228 tracing::warn!("Failed to fetch secrets for {}: {}", provider_id, e);
229 Ok(None)
230 }
231 }
232 }
233
234 pub async fn set_provider_secrets(
236 &self,
237 provider_id: &str,
238 secrets: &ProviderSecrets,
239 ) -> Result<()> {
240 let client = match &self.client {
241 Some(c) => c,
242 None => anyhow::bail!("Vault client not configured"),
243 };
244
245 let secret_path = format!("{}/{}", self.path, provider_id);
246 kv2::set(client.as_ref(), &self.mount, &secret_path, secrets)
247 .await
248 .with_context(|| format!("Failed to write provider secrets for {}", provider_id))?;
249
250 let mut cache = self.cache.write().await;
252 if let Some(api_key) = secrets.api_key.clone() {
253 cache.insert(provider_id.to_string(), api_key);
254 } else {
255 cache.remove(provider_id);
256 }
257
258 Ok(())
259 }
260
261 pub async fn has_api_key(&self, provider_id: &str) -> bool {
263 matches!(self.get_api_key(provider_id).await, Ok(Some(_)))
264 }
265
266 pub async fn list_configured_providers(&self) -> Result<Vec<String>> {
268 let client = match &self.client {
269 Some(c) => c,
270 None => return Ok(Vec::new()),
271 };
272
273 match kv2::list(client.as_ref(), &self.mount, &self.path).await {
274 Ok(keys) => Ok(keys),
275 Err(vaultrs::error::ClientError::APIError { code: 404, .. }) => Ok(Vec::new()),
276 Err(e) => {
277 tracing::warn!("Failed to list providers: {}", e);
278 Ok(Vec::new())
279 }
280 }
281 }
282
283 pub async fn clear_cache(&self) {
285 let mut cache = self.cache.write().await;
286 cache.clear();
287 }
288}
289
290#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
292pub struct VaultConfig {
293 pub address: String,
295
296 pub token: String,
298
299 #[serde(default)]
301 pub mount: Option<String>,
302
303 #[serde(default)]
305 pub path: Option<String>,
306}
307
308impl Default for VaultConfig {
309 fn default() -> Self {
310 Self {
311 address: String::new(),
312 token: String::new(),
313 mount: Some("secret".to_string()),
314 path: Some("codetether/providers".to_string()),
315 }
316 }
317}
318
319#[derive(Clone, serde::Serialize, serde::Deserialize)]
321pub struct ProviderSecrets {
322 #[serde(default)]
324 pub api_key: Option<String>,
325
326 #[serde(default)]
328 pub base_url: Option<String>,
329
330 #[serde(default)]
332 pub organization: Option<String>,
333
334 #[serde(default)]
336 pub headers: Option<HashMap<String, String>>,
337
338 #[serde(flatten)]
340 pub extra: HashMap<String, serde_json::Value>,
341}
342
343impl std::fmt::Debug for ProviderSecrets {
344 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
345 f.debug_struct("ProviderSecrets")
346 .field("api_key", &self.api_key.as_ref().map(|_| "<REDACTED>"))
347 .field("api_key_len", &self.api_key.as_ref().map(|k| k.len()))
348 .field("base_url", &self.base_url)
349 .field("organization", &self.organization)
350 .field("headers_present", &self.headers.is_some())
351 .field("extra_fields", &self.extra.len())
352 .finish()
353 }
354}
355
356impl ProviderSecrets {
357 pub fn has_valid_api_key(&self) -> bool {
359 self.api_key
360 .as_ref()
361 .map(|k| !k.is_empty())
362 .unwrap_or(false)
363 }
364
365 pub fn api_key_len(&self) -> Option<usize> {
367 self.api_key.as_ref().map(|k| k.len())
368 }
369}
370
371static SECRETS_MANAGER: tokio::sync::OnceCell<SecretsManager> = tokio::sync::OnceCell::const_new();
373
374pub async fn init_secrets_manager(config: &VaultConfig) -> Result<()> {
376 let manager = SecretsManager::new(config).await?;
377 SECRETS_MANAGER
378 .set(manager)
379 .map_err(|_| anyhow::anyhow!("Secrets manager already initialized"))?;
380 Ok(())
381}
382
383pub fn init_from_manager(manager: SecretsManager) -> Result<()> {
385 SECRETS_MANAGER
386 .set(manager)
387 .map_err(|_| anyhow::anyhow!("Secrets manager already initialized"))?;
388 Ok(())
389}
390
391pub fn secrets_manager() -> Option<&'static SecretsManager> {
393 SECRETS_MANAGER.get()
394}
395
396pub async fn get_api_key(provider_id: &str) -> Option<String> {
398 match SECRETS_MANAGER.get() {
399 Some(manager) => manager.get_api_key(provider_id).await.ok().flatten(),
400 None => None,
401 }
402}
403
404pub async fn has_api_key(provider_id: &str) -> bool {
406 match SECRETS_MANAGER.get() {
407 Some(manager) => manager.has_api_key(provider_id).await,
408 None => false,
409 }
410}
411
412pub async fn get_provider_secrets(provider_id: &str) -> Option<ProviderSecrets> {
414 match SECRETS_MANAGER.get() {
415 Some(manager) => manager
416 .get_provider_secrets(provider_id)
417 .await
418 .ok()
419 .flatten(),
420 None => None,
421 }
422}
423
424pub async fn set_provider_secrets(provider_id: &str, secrets: &ProviderSecrets) -> Result<()> {
426 match SECRETS_MANAGER.get() {
427 Some(manager) => manager.set_provider_secrets(provider_id, secrets).await,
428 None => anyhow::bail!("Secrets manager not initialized"),
429 }
430}