main
py 559 lines 21.1 KB
Raw
1 import re
2 import threading
3 import time
4 import os
5 from io import StringIO
6 from dataclasses import dataclass
7 from typing import Dict, Optional, List, Literal, Set, Callable, Tuple, TYPE_CHECKING
8 from dotenv.parser import parse_stream
9 from helpers.errors import RepairableException
10 from helpers import dotenv, files
11 from helpers.extension import extensible
12
13 if TYPE_CHECKING:
14 from agent import AgentContext
15
16
17 # New alias-based placeholder format §§secret(KEY)
18 ALIAS_PATTERN = r"§§secret\(([A-Za-z_][A-Za-z0-9_]*)\)"
19 DEFAULT_SECRETS_FILE = "usr/secrets.env"
20 _RUNTIME_CREDENTIAL_KEYS = {
21 dotenv.KEY_AUTH_LOGIN,
22 dotenv.KEY_AUTH_PASSWORD,
23 dotenv.KEY_RFC_PASSWORD,
24 dotenv.KEY_ROOT_PASSWORD,
25 }
26
27
28 def alias_for_key(key: str, placeholder: str = "§§secret({key})") -> str:
29 # Return alias string for given key in upper-case
30 key = key.upper()
31 return placeholder.format(key=key)
32
33
34 @dataclass
35 class EnvLine:
36 raw: str
37 type: Literal["pair", "comment", "blank", "other"]
38 key: Optional[str] = None
39 value: Optional[str] = None
40 inline_comment: Optional[str] = (
41 None # preserves trailing inline comment including leading spaces and '#'
42 )
43
44
45 class StreamingSecretsFilter:
46 """Stateful streaming filter that masks secrets on the fly.
47
48 - Replaces full secret values with placeholders §§secret(KEY) when detected.
49 - Holds the longest suffix of the current buffer that matches any secret prefix
50 (with minimum trigger length of 3) to avoid leaking partial secrets across chunks.
51 - On finalize(), any unresolved partial is masked with '***'.
52 """
53
54 def __init__(self, key_to_value: Dict[str, str], min_trigger: int = 3):
55 self.min_trigger = max(1, int(min_trigger))
56 # Map value -> key for placeholder construction
57 self.value_to_key: Dict[str, str] = {
58 v: k for k, v in key_to_value.items() if isinstance(v, str) and v
59 }
60 # Only keep non-empty values
61 self.secret_values: List[str] = [v for v in self.value_to_key.keys() if v]
62 # Precompute all prefixes for quick suffix matching
63 self.prefixes: Set[str] = set()
64 for v in self.secret_values:
65 for i in range(self.min_trigger, len(v) + 1):
66 self.prefixes.add(v[:i])
67 self.max_len: int = max((len(v) for v in self.secret_values), default=0)
68
69 # Internal buffer of pending text that is not safe to flush yet
70 self.pending: str = ""
71
72 def _replace_full_values(self, text: str) -> str:
73 """Replace all full secret values with placeholders in the given text."""
74 # Sort by length desc to avoid partial overlaps
75 for val in sorted(self.secret_values, key=len, reverse=True):
76 if not val:
77 continue
78 key = self.value_to_key.get(val, "")
79 if key:
80 text = text.replace(val, alias_for_key(key))
81 return text
82
83 def _longest_suffix_prefix(self, text: str) -> int:
84 """Return length of longest suffix of text that is a known secret prefix.
85 Returns 0 if none found (or only shorter than min_trigger)."""
86 max_check = min(len(text), self.max_len)
87 for length in range(max_check, self.min_trigger - 1, -1):
88 suffix = text[-length:]
89 if suffix in self.prefixes:
90 return length
91 return 0
92
93 def process_chunk(self, chunk: str) -> str:
94 if not chunk:
95 return ""
96
97 self.pending += chunk
98
99 # Replace any full secret occurrences first
100 self.pending = self._replace_full_values(self.pending)
101
102 # Determine the longest suffix that could still form a secret
103 hold_len = self._longest_suffix_prefix(self.pending)
104 if hold_len > 0:
105 # Flush everything except the hold suffix
106 emit = self.pending[:-hold_len]
107 self.pending = self.pending[-hold_len:]
108 else:
109 # Safe to flush everything
110 emit = self.pending
111 self.pending = ""
112
113 return emit
114
115 def finalize(self) -> str:
116 """Flush any remaining buffered text. If pending contains an unresolved partial
117 (i.e., a prefix of a secret >= min_trigger), mask it with *** to avoid leaks."""
118 if not self.pending:
119 return ""
120
121 hold_len = self._longest_suffix_prefix(self.pending)
122 if hold_len > 0:
123 safe = self.pending[:-hold_len]
124 # Mask unresolved partial
125 result = safe + "***"
126 else:
127 result = self.pending
128 self.pending = ""
129 return result
130
131
132 class SecretsManager:
133 PLACEHOLDER_PATTERN = ALIAS_PATTERN
134 MASK_VALUE = "***"
135
136 _instances: Dict[Tuple[str, ...], "SecretsManager"] = {}
137 _secrets_cache: Optional[Dict[str, str]] = None
138 _last_raw_text: Optional[str] = None
139
140 @classmethod
141 def get_instance(cls, *secrets_files: str) -> "SecretsManager":
142 if not secrets_files:
143 secrets_files = (DEFAULT_SECRETS_FILE,)
144 key = tuple(secrets_files)
145 if key not in cls._instances:
146 cls._instances[key] = cls(*secrets_files)
147 return cls._instances[key]
148
149 def __init__(self, *files: str):
150 self._lock = threading.RLock()
151 # instance-level list of secrets files
152 self._files: Tuple[str, ...] = tuple(files) if files else (DEFAULT_SECRETS_FILE,)
153 self._raw_snapshots: Dict[str, str] = {}
154 self._secrets_cache = None
155 self._last_raw_text = None
156
157 def read_secrets_raw(self) -> str:
158 """Read raw secrets file content from local filesystem (same system)."""
159 parts: List[str] = []
160 self._raw_snapshots = {}
161
162 for path in self._files:
163 try:
164 content = files.read_file(path)
165 except Exception:
166 content = ""
167
168 self._raw_snapshots[path] = content
169 if path == dotenv.get_dotenv_file_path():
170 content = "\n".join(
171 line.raw
172 for line in self.parse_env_lines(content)
173 if line.type != "pair"
174 or (line.key or "").upper().startswith("API_KEY_")
175 or (line.key or "").upper() in _RUNTIME_CREDENTIAL_KEYS
176 )
177 parts.append(content)
178
179 combined = "\n".join(parts)
180 self._last_raw_text = combined
181 return combined
182
183 def _write_secrets_raw(self, content: str):
184 """Write raw secrets file content to local filesystem."""
185 if len(self._files) != 1:
186 raise RuntimeError(
187 "Saving secrets content is only supported for a single secrets file"
188 )
189 files.write_file(self._files[0], content)
190
191 def load_secrets(self) -> Dict[str, str]:
192 """Load secrets from file, return key-value dict"""
193 with self._lock:
194 if self._secrets_cache is not None:
195 return self._secrets_cache
196
197 combined_raw = self.read_secrets_raw()
198 merged_secrets = (
199 self.parse_env_content(combined_raw) if combined_raw else {}
200 )
201
202 # Only track the first file's raw text for single-file setups
203 if len(self._files) != 1:
204 self._last_raw_text = None
205
206 self._secrets_cache = merged_secrets
207 return merged_secrets
208
209 def save_secrets(self, secrets_content: str):
210 """Save secrets content to file and update cache"""
211 if len(self._files) != 1:
212 raise RuntimeError(
213 "Saving secrets is disabled when multiple files are configured"
214 )
215 with self._lock:
216 self._write_secrets_raw(secrets_content)
217 self._invalidate_all_caches()
218
219 def save_secrets_with_merge(self, submitted_content: str):
220 """Merge submitted content with existing file preserving comments, order and supporting deletion.
221 - Existing keys keep their value when submitted as MASK_VALUE (***).
222 - Keys present in existing but omitted from submitted are deleted.
223 - New keys with non-masked values are appended at the end.
224 """
225 if len(self._files) != 1:
226 raise RuntimeError(
227 "Merging secrets is disabled when multiple files are configured"
228 )
229 with self._lock:
230 # Prefer in-memory snapshot to avoid disk reads during save
231 primary_path = self._files[0]
232 if self._last_raw_text is not None:
233 existing_text = self._last_raw_text
234 else:
235 try:
236 existing_text = files.read_file(primary_path)
237 self._raw_snapshots[primary_path] = existing_text
238 except Exception as e:
239 # If read fails and submitted contains masked values, abort to avoid losing values/comments
240 if self.MASK_VALUE in submitted_content:
241 raise RepairableException(
242 "Saving secrets failed because existing secrets could not be read to preserve masked values and comments. Please retry."
243 ) from e
244 # No masked values, safe to treat as new file
245 existing_text = ""
246 merged_lines = self._merge_env(existing_text, submitted_content)
247 merged_text = self._serialize_env_lines(merged_lines)
248 self._write_secrets_raw(merged_text)
249 self._invalidate_all_caches()
250
251 def get_keys(self) -> List[str]:
252 """Get list of secret keys"""
253 secrets = self.load_secrets()
254 return list(secrets.keys())
255
256 def get_secrets_for_prompt(self) -> str:
257 """Get formatted string of secret keys for system prompt"""
258 content = self.read_secrets_raw()
259 if not content:
260 return ""
261
262 env_lines = self.parse_env_lines(content)
263 return self._serialize_env_lines(
264 env_lines,
265 with_values=False,
266 with_comments=True,
267 with_blank=True,
268 with_other=True,
269 key_formatter=alias_for_key,
270 )
271
272 def create_streaming_filter(self) -> "StreamingSecretsFilter":
273 """Create a streaming-aware secrets filter snapshotting current secret values."""
274 return StreamingSecretsFilter(self.load_secrets())
275
276 def replace_placeholders(self, text: str) -> str:
277 """Replace secret placeholders with actual values"""
278 if not text:
279 return text
280
281 secrets = self.load_secrets()
282
283 def replacer(match):
284 key = match.group(1)
285 key = key.upper()
286 if key in secrets:
287 return secrets[key]
288 else:
289 available_keys = ", ".join(secrets.keys())
290 error_msg = f"Secret placeholder '{alias_for_key(key)}' not found in secrets store.\n"
291 error_msg += f"Available secrets: {available_keys}"
292
293 raise RepairableException(error_msg)
294
295 return re.sub(self.PLACEHOLDER_PATTERN, replacer, text)
296
297 def change_placeholders(self, text: str, new_format: str) -> str:
298 """Substitute secret placeholders with a different placeholder format"""
299 if not text:
300 return text
301
302 secrets = self.load_secrets()
303 result = text
304
305 # Sort by length (longest first) to avoid partial replacements
306 for key, _value in sorted(
307 secrets.items(), key=lambda x: len(x[1]), reverse=True
308 ):
309 result = result.replace(alias_for_key(key), new_format.format(key=key))
310
311 return result
312
313 def mask_values(
314 self, text: str, min_length: int = 4, placeholder: str = "§§secret({key})"
315 ) -> str:
316 """Replace actual secret values with placeholders in text"""
317 if not text:
318 return text
319
320 secrets = self.load_secrets()
321 result = text
322
323 # Sort by length (longest first) to avoid partial replacements
324 for key, value in sorted(
325 secrets.items(), key=lambda x: len(x[1]), reverse=True
326 ):
327 if value and len(value.strip()) >= min_length:
328 result = result.replace(value, alias_for_key(key, placeholder))
329
330 return result
331
332 def get_masked_secrets(self) -> str:
333 """Get content with values masked for frontend display (preserves comments and unrecognized lines)"""
334 content = self.read_secrets_raw()
335 if not content:
336 return ""
337
338 # Parse content for known keys using python-dotenv
339 secrets_map = self.parse_env_content(content)
340 env_lines = self.parse_env_lines(content)
341
342 # Replace values with mask for keys present
343 for ln in env_lines:
344 if ln.type == "pair" and ln.key is not None:
345 ln.key = ln.key.upper()
346 if ln.key in secrets_map and secrets_map[ln.key] != "":
347 ln.value = self.MASK_VALUE
348
349 return self._serialize_env_lines(env_lines)
350
351 def parse_env_content(self, content: str) -> Dict[str, str]:
352 """Parse .env format content into key-value dict using python-dotenv. Keys are always uppercase."""
353 env: Dict[str, str] = {}
354 for binding in parse_stream(StringIO(content)):
355 if binding.key and not binding.error:
356 env[binding.key.upper()] = binding.value or ""
357 return env
358
359 # Backward-compatible alias for callers using the old private method name
360 def _parse_env_content(self, content: str) -> Dict[str, str]:
361 return self.parse_env_content(content)
362
363 def clear_cache(self):
364 """Clear the secrets cache"""
365 with self._lock:
366 self._secrets_cache = None
367 self._raw_snapshots = {}
368 self._last_raw_text = None
369
370 @classmethod
371 def _invalidate_all_caches(cls):
372 for instance in cls._instances.values():
373 instance.clear_cache()
374
375 # ---------------- Internal helpers for parsing/merging ----------------
376
377 def parse_env_lines(self, content: str) -> List[EnvLine]:
378 """Parse env file into EnvLine objects using python-dotenv, preserving comments and order.
379 We reconstruct key_part and inline_comment based on the original string.
380 """
381 lines: List[EnvLine] = []
382 for binding in parse_stream(StringIO(content)):
383 orig = getattr(binding, "original", None)
384 raw = getattr(orig, "string", "") if orig is not None else ""
385 if binding.key and not binding.error:
386 # Determine key_part and inline_comment from original line
387 line_text = raw.rstrip("\n")
388 # Fallback to composed key_part if original not available
389 if "=" in line_text:
390 left, right = line_text.split("=", 1)
391 else:
392 right = ""
393 # Try to extract inline comment by scanning right side to comment start, respecting quotes
394 in_single = False
395 in_double = False
396 esc = False
397 comment_index = None
398 for i, ch in enumerate(right):
399 if esc:
400 esc = False
401 continue
402 if ch == "\\":
403 esc = True
404 continue
405 if ch == "'" and not in_double:
406 in_single = not in_single
407 continue
408 if ch == '"' and not in_single:
409 in_double = not in_double
410 continue
411 if ch == "#" and not in_single and not in_double:
412 comment_index = i
413 break
414 inline_comment = None
415 if comment_index is not None:
416 inline_comment = right[comment_index:]
417 lines.append(
418 EnvLine(
419 raw=line_text,
420 type="pair",
421 key=binding.key,
422 value=binding.value or "",
423 inline_comment=inline_comment,
424 )
425 )
426 else:
427 # Comment, blank, or other lines
428 raw_line = raw.rstrip("\n")
429 if raw_line.strip() == "":
430 lines.append(EnvLine(raw=raw_line, type="blank"))
431 elif raw_line.lstrip().startswith("#"):
432 lines.append(EnvLine(raw=raw_line, type="comment"))
433 else:
434 lines.append(EnvLine(raw=raw_line, type="other"))
435 return lines
436
437 def _serialize_env_lines(
438 self,
439 lines: List[EnvLine],
440 with_values=True,
441 with_comments=True,
442 with_blank=True,
443 with_other=True,
444 key_delimiter="",
445 key_formatter: Optional[Callable[[str], str]] = None,
446 ) -> str:
447 out: List[str] = []
448 for ln in lines:
449 if ln.type == "pair" and ln.key is not None:
450 left_raw = ln.key
451 left = left_raw.upper()
452 val = ln.value if ln.value is not None else ""
453 comment = ln.inline_comment or ""
454 formatted_key = (
455 key_formatter(left)
456 if key_formatter
457 else f"{key_delimiter}{left}{key_delimiter}"
458 )
459 val_part = f'="{val}"' if with_values else ""
460 comment_part = f" {comment}" if with_comments and comment else ""
461 out.append(f"{formatted_key}{val_part}{comment_part}")
462 elif ln.type == "blank" and with_blank:
463 out.append(ln.raw)
464 elif ln.type == "comment" and with_comments:
465 out.append(ln.raw)
466 elif ln.type == "other" and with_other:
467 out.append(ln.raw)
468 return "\n".join(out)
469
470 def _merge_env(self, existing_text: str, submitted_text: str) -> List[EnvLine]:
471 """Merge using submitted content as the base to preserve its comments and structure.
472 Behavior:
473 - Iterate submitted lines in order and keep them (including comments/blanks/other).
474 - For pair lines:
475 - If key exists in existing and submitted value is MASK_VALUE (***), use existing value.
476 - If key is new and value is MASK_VALUE, skip (ignore masked-only additions).
477 - Otherwise, use submitted value as-is.
478 - Keys present only in existing and not in submitted are deleted (not added).
479 This preserves comments and arbitrary lines from the submitted content and persists them.
480 """
481 existing_lines = self.parse_env_lines(existing_text)
482 submitted_lines = self.parse_env_lines(submitted_text)
483
484 existing_pairs: Dict[str, EnvLine] = {
485 ln.key: ln
486 for ln in existing_lines
487 if ln.type == "pair" and ln.key is not None
488 }
489
490 merged: List[EnvLine] = []
491 for sub in submitted_lines:
492 if sub.type != "pair" or sub.key is None:
493 # Preserve submitted comments/blanks/other verbatim
494 merged.append(sub)
495 continue
496
497 key = sub.key
498 submitted_val = sub.value or ""
499
500 if key in existing_pairs and submitted_val == self.MASK_VALUE:
501 # Replace mask with existing value, keep submitted key formatting
502 existing_val = existing_pairs[key].value or ""
503 merged.append(
504 EnvLine(
505 raw=f"{key}={existing_val}",
506 type="pair",
507 key=key,
508 value=existing_val,
509 inline_comment=sub.inline_comment,
510 )
511 )
512 elif key not in existing_pairs and submitted_val == self.MASK_VALUE:
513 # Masked-only new key -> ignore
514 continue
515 else:
516 # Use submitted value as-is
517 merged.append(sub)
518
519 return merged
520
521
522 @extensible
523 def get_secrets_manager(context: "AgentContext|None" = None) -> SecretsManager:
524 from helpers import projects
525
526 # Agent-facing masking covers the secret store and runtime credentials in usr/.env.
527 secret_files = [DEFAULT_SECRETS_FILE, dotenv.get_dotenv_file_path()]
528
529 # use AgentContext from contextvars if no context provided
530 if not context:
531 from agent import AgentContext
532 context = AgentContext.current()
533
534 # merged with project secrets if active
535 if context:
536 project = projects.get_context_project_name(context)
537 if project:
538 secret_files.append(files.get_abs_path(projects.get_project_meta(project), "secrets.env"))
539
540 return SecretsManager.get_instance(*secret_files)
541
542 @extensible
543 def get_project_secrets_manager(project_name: str, merge_with_global: bool = False) -> SecretsManager:
544 from helpers import projects
545
546 # default secrets file
547 secret_files = []
548
549 if merge_with_global:
550 secret_files.append(DEFAULT_SECRETS_FILE)
551
552 # merged with project secrets if active
553 secret_files.append(files.get_abs_path(projects.get_project_meta(project_name), "secrets.env"))
554
555 return SecretsManager.get_instance(*secret_files)
556
557 @extensible
558 def get_default_secrets_manager() -> SecretsManager:
559 return SecretsManager.get_instance()