Coverage for agentos/core/secrets.py: 72%

1"""AgentOS Secrets Manager — production-grade secrets lifecycle. 

2 

3Backends: 

4- EnvSecretsBackend: os.environ (12-factor) 

5- VaultSecretsBackend: HashiCorp Vault (kv-v2) 

6- EncryptedFileBackend: Fernet-encrypted JSON file 

7- CompositeSecretsBackend: layered resolution (most specific first) 

8 

9Design: ~350 lines, thread-safe, lazy loading with caching. 

10""" 

11 

12from __future__ import annotations 

13 

14import json 

15import logging 

16import os 

17import time 

18import warnings 

19from abc import ABC, abstractmethod 

20from dataclasses import dataclass, field 

21from enum import Enum 

22from pathlib import Path 

23from threading import RLock 

24from typing import Any, Callable, Dict, List, Optional 

25 

26logger = logging.getLogger(__name__) 

27 

28 

29# ============================================================================ 

30# Data types 

31# ============================================================================ 

32 

33class SecretNotFoundError(Exception): 

34 """Requested secret key does not exist.""" 

35 

36 

37class BackendUnavailableError(Exception): 

38 """Secrets backend is unreachable or misconfigured.""" 

39 

40 

41@dataclass 

42class SecretsConfig: 

43 """Global secrets manager configuration.""" 

44 

45 cache_ttl: float = 300.0 # Seconds to cache fetched secrets 

46 max_cache_size: int = 1000 # Max cached entries 

47 fail_open: bool = False # If True, return None on backend error instead of raising 

48 allow_environment_fallback: bool = True # Try env vars before backend 

49 

50 

51# ============================================================================ 

52# Abstract backend 

53# ============================================================================ 

54 

55class AbstractSecretsBackend(ABC): 

56 """Interface for secrets backends.""" 

57 

58 @abstractmethod 

59 async def get(self, key: str) -> Optional[str]: 

60 """Retrieve a single secret value.""" 

61 

62 @abstractmethod 

63 async def get_all(self, prefix: str = "") -> Dict[str, str]: 

64 """Retrieve all secrets matching prefix.""" 

65 

66 @abstractmethod 

67 async def health_check(self) -> bool: 

68 """Verify backend connectivity.""" 

69 

70 

71# ============================================================================ 

72# Env backend 

73# ============================================================================ 

74 

75class EnvSecretsBackend(AbstractSecretsBackend): 

76 """OS environment variable backend — no deps, always available.""" 

77 

78 def __init__(self, prefix: str = ""): 

79 self._prefix = prefix 

80 

81 async def get(self, key: str) -> Optional[str]: 

82 return os.environ.get(f"{self._prefix}{key}") 

83 

84 async def get_all(self, prefix: str = "") -> Dict[str, str]: 

85 full_prefix = f"{self._prefix}{prefix}" 

86 return { 

87 k[len(self._prefix):] if k.startswith(self._prefix) else k: v 

88 for k, v in os.environ.items() 

89 if k.startswith(full_prefix) 

90 } 

91 

92 async def health_check(self) -> bool: 

93 return True 

94 

95 

96# ============================================================================ 

97# Encrypted file backend 

98# ============================================================================ 

99 

100class EncryptedFileBackend(AbstractSecretsBackend): 

101 """Fernet-encrypted JSON file backend. 

102 

103 Secrets file encrypted at rest. Use `cryptography` for encryption. 

104 """ 

105 

106 def __init__(self, file_path: str, encryption_key: str): 

107 self._path = Path(file_path) 

108 self._key = encryption_key 

109 self._cache: Optional[Dict[str, str]] = None 

110 self._cache_time: float = 0.0 

111 self._lock = RLock() 

112 

113 def _decrypt(self) -> Dict[str, str]: 

114 """Decrypt and load secrets file.""" 

115 from cryptography.fernet import Fernet 

116 

117 if not self._path.exists(): 

118 return {} 

119 

120 fernet = Fernet(self._key.encode() if isinstance(self._key, str) else self._key) 

121 with open(self._path, "rb") as f: 

122 encrypted = f.read() 

123 decrypted = fernet.decrypt(encrypted) 

124 return json.loads(decrypted) 

125 

126 async def get(self, key: str) -> Optional[str]: 

127 with self._lock: 

128 if self._cache is None or (time.monotonic() - self._cache_time) > 300: 

129 try: 

130 self._cache = self._decrypt() 

131 self._cache_time = time.monotonic() 

132 except Exception as exc: 

133 logger.error("Failed to decrypt secrets file: %s", exc) 

134 raise BackendUnavailableError(f"Decryption failed: {exc}") 

135 return self._cache.get(key) if self._cache else None 

136 

137 async def get_all(self, prefix: str = "") -> Dict[str, str]: 

138 await self.get("") # Refresh cache 

139 if self._cache is None: 

140 return {} 

141 if not prefix: 

142 return dict(self._cache) 

143 return {k: v for k, v in self._cache.items() if k.startswith(prefix)} 

144 

145 async def health_check(self) -> bool: 

146 try: 

147 self._decrypt() 

148 return True 

149 except Exception: 

150 return False 

151 

152 

153# ============================================================================ 

154# Vault backend (HashiCorp Vault kv-v2) 

155# ============================================================================ 

156 

157class VaultSecretsBackend(AbstractSecretsBackend): 

158 """HashiCorp Vault backend (kv-v2). 

159 

160 Requires hvac library or httpx for REST calls. 

161 """ 

162 

163 def __init__( 

164 self, 

165 url: str, 

166 token: str, 

167 mount_point: str = "secret", 

168 path_prefix: str = "", 

169 verify_ssl: bool = True, 

170 ): 

171 self._url = url.rstrip("/") 

172 self._token = token 

173 self._mount_point = mount_point 

174 self._path_prefix = path_prefix 

175 self._verify_ssl = verify_ssl 

176 self._cache: Dict[str, Optional[str]] = {} 

177 self._lock = RLock() 

178 

179 async def _call(self, method: str, path: str) -> Any: 

180 """Make authenticated Vault API call.""" 

181 import httpx 

182 

183 async with httpx.AsyncClient(verify=self._verify_ssl, timeout=10.0) as client: 

184 headers = {"X-Vault-Token": self._token} 

185 url = f"{self._url}/v1/{path}" 

186 resp = await client.request(method, url, headers=headers) 

187 if resp.status_code == 404: 

188 return None 

189 if resp.status_code >= 400: 

190 raise BackendUnavailableError( 

191 f"Vault API error {resp.status_code}: {resp.text[:200]}" 

192 ) 

193 return resp.json() 

194 

195 async def get(self, key: str) -> Optional[str]: 

196 full_path = f"{self._path_prefix}/{key}" if self._path_prefix else key 

197 vault_path = f"{self._mount_point}/data/{full_path}" 

198 

199 with self._lock: 

200 if key in self._cache: 

201 return self._cache[key] 

202 

203 try: 

204 data = await self._call("GET", vault_path) 

205 if data is None: 

206 return None 

207 value = data.get("data", {}).get("data", {}).get("value") 

208 if value is None: 

209 value = data.get("data", {}).get("data", {}) 

210 with self._lock: 

211 self._cache[key] = str(value) if not isinstance(value, dict) else json.dumps(value) 

212 return self._cache[key] 

213 except BackendUnavailableError: 

214 raise 

215 

216 async def get_all(self, prefix: str = "") -> Dict[str, str]: 

217 list_path = f"{self._mount_point}/metadata/{self._path_prefix}" 

218 result: Dict[str, str] = {} 

219 

220 try: 

221 data = await self._call("LIST", list_path) 

222 if data and "data" in data and "keys" in data["data"]: 

223 for key in data["data"]["keys"]: 

224 if not prefix or key.startswith(prefix): 

225 value = await self.get(key) 

226 if value is not None: 

227 result[key] = value 

228 except BackendUnavailableError: 

229 raise 

230 

231 return result 

232 

233 async def health_check(self) -> bool: 

234 try: 

235 resp = await self._call("GET", "sys/health") 

236 return resp is not None and resp.get("initialized", False) 

237 except Exception: 

238 return False 

239 

240 

241# ============================================================================ 

242# Composite backend (layered) 

243# ============================================================================ 

244 

245class CompositeSecretsBackend(AbstractSecretsBackend): 

246 """Resolve secrets from multiple backends in priority order. 

247 

248 First backend that returns a non-None value wins. 

249 """ 

250 

251 def __init__(self, backends: List[AbstractSecretsBackend]): 

252 self._backends = backends 

253 

254 async def get(self, key: str) -> Optional[str]: 

255 for backend in self._backends: 

256 try: 

257 value = await backend.get(key) 

258 if value is not None: 

259 return value 

260 except Exception as exc: 

261 logger.debug("Backend %s failed for key=%s: %s", type(backend).__name__, key, exc) 

262 return None 

263 

264 async def get_all(self, prefix: str = "") -> Dict[str, str]: 

265 result: Dict[str, str] = {} 

266 for backend in reversed(self._backends): # Low-priority first 

267 try: 

268 batch = await backend.get_all(prefix) 

269 result.update(batch) 

270 except Exception as exc: 

271 logger.debug("Backend %s get_all failed: %s", type(backend).__name__, exc) 

272 return result 

273 

274 async def health_check(self) -> bool: 

275 for backend in self._backends: 

276 if await backend.health_check(): 

277 return True 

278 return False 

279 

280 

281# ============================================================================ 

282# Secrets Manager (high-level) 

283# ============================================================================ 

284 

285class SecretsManager: 

286 """High-level secrets manager with caching and fail-open support. 

287 

288 Usage: 

289 sm = SecretsManager(EnvSecretsBackend("MYAPP_")) 

290 api_key = await sm.get("API_KEY") 

291 # or 

292 api_key = await sm.require("API_KEY") # raises if missing 

293 """ 

294 

295 def __init__(self, backend: AbstractSecretsBackend, config: SecretsConfig = SecretsConfig()): 

296 self._backend = backend 

297 self._config = config 

298 self._cache: Dict[str, tuple[float, Optional[str]]] = {} 

299 self._lock = RLock() 

300 

301 async def _cache_put(self, key: str, value: Optional[str]): 

302 """Store value in cache with eviction.""" 

303 with self._lock: 

304 if len(self._cache) >= self._config.max_cache_size: 

305 sorted_keys = sorted(self._cache, key=lambda k: self._cache[k][0]) 

306 for old_key in sorted_keys[:len(self._cache) // 4]: 

307 del self._cache[old_key] 

308 self._cache[key] = (time.monotonic(), value) 

309 

310 async def get(self, key: str) -> Optional[str]: 

311 """Get secret value. Returns None if not found.""" 

312 # Check cache 

313 with self._lock: 

314 if key in self._cache: 

315 ts, val = self._cache[key] 

316 if time.monotonic() - ts < self._config.cache_ttl: 

317 return val 

318 

319 # Fallback to environment 

320 if self._config.allow_environment_fallback: 

321 env_val = os.environ.get(key) 

322 if env_val is not None: 

323 await self._cache_put(key, env_val) 

324 return env_val 

325 

326 # Query backend 

327 try: 

328 value = await self._backend.get(key) 

329 except Exception as exc: 

330 if self._config.fail_open: 

331 logger.warning("Secrets backend error for key=%s (fail_open): %s", key, exc) 

332 return None 

333 raise BackendUnavailableError(f"Failed to fetch '{key}': {exc}") from exc 

334 

335 # Update cache 

336 await self._cache_put(key, value) 

337 

338 return value 

339 

340 async def require(self, key: str) -> str: 

341 """Get secret value. Raises SecretNotFoundError if missing.""" 

342 value = await self.get(key) 

343 if value is None: 

344 raise SecretNotFoundError(f"Required secret '{key}' not found") 

345 return value 

346 

347 async def get_all(self, prefix: str = "") -> Dict[str, str]: 

348 """Get all secrets with given prefix.""" 

349 try: 

350 return await self._backend.get_all(prefix) 

351 except Exception as exc: 

352 if self._config.fail_open: 

353 logger.warning("Secrets get_all error (fail_open): %s", exc) 

354 return {} 

355 raise 

356 

357 async def health_check(self) -> bool: 

358 """Check backend health.""" 

359 return await self._backend.health_check() 

360 

361 def invalidate_cache(self, key: Optional[str] = None): 

362 """Invalidate cache entries.""" 

363 with self._lock: 

364 if key is None: 

365 self._cache.clear() 

366 elif key in self._cache: 

367 del self._cache[key] 

368 

369 

370# ============================================================================ 

371# Convenience factory 

372# ============================================================================ 

373 

374def create_secrets_manager( 

375 backend_type: str = "env", 

376 **kwargs: Any, 

377) -> SecretsManager: 

378 """Factory: create SecretsManager with common backends. 

379 

380 backend_type: 'env' | 'vault' | 'encrypted_file' | 'composite' 

381 """ 

382 if backend_type == "env": 

383 backend = EnvSecretsBackend(prefix=kwargs.get("prefix", "")) 

384 elif backend_type == "vault": 

385 backend = VaultSecretsBackend( 

386 url=kwargs["url"], 

387 token=kwargs["token"], 

388 mount_point=kwargs.get("mount_point", "secret"), 

389 path_prefix=kwargs.get("path_prefix", ""), 

390 verify_ssl=kwargs.get("verify_ssl", True), 

391 ) 

392 elif backend_type == "encrypted_file": 

393 backend = EncryptedFileBackend( 

394 file_path=kwargs["file_path"], 

395 encryption_key=kwargs["encryption_key"], 

396 ) 

397 elif backend_type == "composite": 

398 backends = kwargs["backends"] 

399 backend = CompositeSecretsBackend(backends) 

400 else: 

401 raise ValueError(f"Unknown backend type: {backend_type}") 

402 

403 return SecretsManager(backend, SecretsConfig(**kwargs.get("config", {})))