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

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 

18from abc import ABC, abstractmethod 

19from dataclasses import dataclass 

20from pathlib import Path 

21from threading import RLock 

22from typing import Any, Dict, List, Optional 

23 

24logger = logging.getLogger(__name__) 

25 

26 

27# ============================================================================ 

28# Data types 

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

30 

31class SecretNotFoundError(Exception): 

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

33 

34 

35class BackendUnavailableError(Exception): 

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

37 

38 

39@dataclass 

40class SecretsConfig: 

41 """Global secrets manager configuration.""" 

42 

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

44 max_cache_size: int = 1000 # Max cached entries 

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

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

47 

48 

49# ============================================================================ 

50# Abstract backend 

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

52 

53class AbstractSecretsBackend(ABC): 

54 """Interface for secrets backends.""" 

55 

56 @abstractmethod 

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

58 """Retrieve a single secret value.""" 

59 

60 @abstractmethod 

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

62 """Retrieve all secrets matching prefix.""" 

63 

64 @abstractmethod 

65 async def health_check(self) -> bool: 

66 """Verify backend connectivity.""" 

67 

68 

69# ============================================================================ 

70# Env backend 

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

72 

73class EnvSecretsBackend(AbstractSecretsBackend): 

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

75 

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

77 self._prefix = prefix 

78 

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

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

81 

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

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

84 return { 

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

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

87 if k.startswith(full_prefix) 

88 } 

89 

90 async def health_check(self) -> bool: 

91 return True 

92 

93 

94# ============================================================================ 

95# Encrypted file backend 

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

97 

98class EncryptedFileBackend(AbstractSecretsBackend): 

99 """Fernet-encrypted JSON file backend. 

100 

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

102 """ 

103 

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

105 self._path = Path(file_path) 

106 self._key = encryption_key 

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

108 self._cache_time: float = 0.0 

109 self._lock = RLock() 

110 

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

112 """Decrypt and load secrets file.""" 

113 from cryptography.fernet import Fernet 

114 

115 if not self._path.exists(): 

116 return {} 

117 

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

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

120 encrypted = f.read() 

121 decrypted = fernet.decrypt(encrypted) 

122 return json.loads(decrypted) 

123 

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

125 with self._lock: 

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

127 try: 

128 self._cache = self._decrypt() 

129 self._cache_time = time.monotonic() 

130 except Exception as exc: 

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

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

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

134 

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

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

137 if self._cache is None: 

138 return {} 

139 if not prefix: 

140 return dict(self._cache) 

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

142 

143 async def health_check(self) -> bool: 

144 try: 

145 self._decrypt() 

146 return True 

147 except Exception: 

148 return False 

149 

150 

151# ============================================================================ 

152# Vault backend (HashiCorp Vault kv-v2) 

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

154 

155class VaultSecretsBackend(AbstractSecretsBackend): 

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

157 

158 Requires hvac library or httpx for REST calls. 

159 """ 

160 

161 def __init__( 

162 self, 

163 url: str, 

164 token: str, 

165 mount_point: str = "secret", 

166 path_prefix: str = "", 

167 verify_ssl: bool = True, 

168 ): 

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

170 self._token = token 

171 self._mount_point = mount_point 

172 self._path_prefix = path_prefix 

173 self._verify_ssl = verify_ssl 

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

175 self._lock = RLock() 

176 

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

178 """Make authenticated Vault API call.""" 

179 import httpx 

180 

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

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

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

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

185 if resp.status_code == 404: 

186 return None 

187 if resp.status_code >= 400: 

188 raise BackendUnavailableError( 

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

190 ) 

191 return resp.json() 

192 

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

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

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

196 

197 with self._lock: 

198 if key in self._cache: 

199 return self._cache[key] 

200 

201 try: 

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

203 if data is None: 

204 return None 

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

206 if value is None: 

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

208 with self._lock: 

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

210 return self._cache[key] 

211 except BackendUnavailableError: 

212 raise 

213 

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

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

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

217 

218 try: 

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

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

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

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

223 value = await self.get(key) 

224 if value is not None: 

225 result[key] = value 

226 except BackendUnavailableError: 

227 raise 

228 

229 return result 

230 

231 async def health_check(self) -> bool: 

232 try: 

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

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

235 except Exception: 

236 return False 

237 

238 

239# ============================================================================ 

240# Composite backend (layered) 

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

242 

243class CompositeSecretsBackend(AbstractSecretsBackend): 

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

245 

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

247 """ 

248 

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

250 self._backends = backends 

251 

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

253 for backend in self._backends: 

254 try: 

255 value = await backend.get(key) 

256 if value is not None: 

257 return value 

258 except Exception as exc: 

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

260 return None 

261 

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

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

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

265 try: 

266 batch = await backend.get_all(prefix) 

267 result.update(batch) 

268 except Exception as exc: 

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

270 return result 

271 

272 async def health_check(self) -> bool: 

273 for backend in self._backends: 

274 if await backend.health_check(): 

275 return True 

276 return False 

277 

278 

279# ============================================================================ 

280# Secrets Manager (high-level) 

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

282 

283class SecretsManager: 

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

285 

286 Usage: 

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

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

289 # or 

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

291 """ 

292 

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

294 self._backend = backend 

295 self._config = config 

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

297 self._lock = RLock() 

298 

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

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

301 with self._lock: 

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

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

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

305 del self._cache[old_key] 

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

307 

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

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

310 # Check cache 

311 with self._lock: 

312 if key in self._cache: 

313 ts, val = self._cache[key] 

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

315 return val 

316 

317 # Fallback to environment 

318 if self._config.allow_environment_fallback: 

319 env_val = os.environ.get(key) 

320 if env_val is not None: 

321 await self._cache_put(key, env_val) 

322 return env_val 

323 

324 # Query backend 

325 try: 

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

327 except Exception as exc: 

328 if self._config.fail_open: 

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

330 return None 

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

332 

333 # Update cache 

334 await self._cache_put(key, value) 

335 

336 return value 

337 

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

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

340 value = await self.get(key) 

341 if value is None: 

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

343 return value 

344 

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

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

347 try: 

348 return await self._backend.get_all(prefix) 

349 except Exception as exc: 

350 if self._config.fail_open: 

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

352 return {} 

353 raise 

354 

355 async def health_check(self) -> bool: 

356 """Check backend health.""" 

357 return await self._backend.health_check() 

358 

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

360 """Invalidate cache entries.""" 

361 with self._lock: 

362 if key is None: 

363 self._cache.clear() 

364 elif key in self._cache: 

365 del self._cache[key] 

366 

367 

368# ============================================================================ 

369# Convenience factory 

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

371 

372def create_secrets_manager( 

373 backend_type: str = "env", 

374 **kwargs: Any, 

375) -> SecretsManager: 

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

377 

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

379 """ 

380 if backend_type == "env": 

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

382 elif backend_type == "vault": 

383 backend = VaultSecretsBackend( 

384 url=kwargs["url"], 

385 token=kwargs["token"], 

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

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

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

389 ) 

390 elif backend_type == "encrypted_file": 

391 backend = EncryptedFileBackend( 

392 file_path=kwargs["file_path"], 

393 encryption_key=kwargs["encryption_key"], 

394 ) 

395 elif backend_type == "composite": 

396 backends = kwargs["backends"] 

397 backend = CompositeSecretsBackend(backends) 

398 else: 

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

400 

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