[core] Add YAML secret reference and registration helpers (#20343)

esphome/bundle.py

@@ -102,25 +102,21 @@ _KNOWN_FILE_EXTENSIONS = frozenset(
 )
 
 
-# Matches !secret references in YAML text.  An optional surrounding
-# quote pair around the key is allowed and ignored: YAML treats
-# ``!secret 'foo'`` and ``!secret foo`` as the same key.  This is
-# intentionally a simple regex scan rather than a YAML parse — it may
-# match inside comments or multi-line strings, which is the conservative
-# direction (include more secrets rather than fewer).
-_SECRET_RE = re.compile(r"""!secret\s+['"]?([^\s'"]+)""")
-
-
 def _find_used_secret_keys(yaml_files: list[Path]) -> set[str]:
     """Scan YAML files for ``!secret <key>`` references."""
     keys: set[str] = set()
     for fpath in yaml_files:
         try:
             text = fpath.read_text(encoding="utf-8")
-        except (OSError, UnicodeDecodeError):
+        except (OSError, UnicodeDecodeError) as err:
+            _LOGGER.warning(
+                "Could not scan %s for !secret references (%s); the bundled "
+                "secret set may be incomplete",
+                fpath,
+                err,
+            )
             continue
-        for match in _SECRET_RE.finditer(text):
-            keys.add(match.group(1))
+        keys |= yaml_util.find_secret_references(text)
     return keys
 
 

esphome/yaml_util.py

@@ -11,6 +11,7 @@ import logging
 import math
 import os
 from pathlib import Path
+import re
 from typing import Any
 import uuid
 
@@ -46,6 +47,12 @@ _LOGGER = logging.getLogger(__name__)
 SECRET_YAML = "secrets.yaml"
 _SECRET_CACHE = {}
 _SECRET_VALUES = {}
+# Every !secret name the loader resolved; _SECRET_VALUES is keyed by value, so
+# names sharing a value collide there.
+_LOADED_SECRET_NAMES: set[str] = set()
+# Set while a secret_values_registered context is active; the dumper records
+# every emitted `!secret` name into it.
+_EMITTED_SECRET_NAMES: set[str] | None = None
 # Not thread-safe — config processing is single-threaded today.
 _load_listeners: list[Callable[[Path], None]] = []
 
@@ -456,6 +463,16 @@ def _load_include_candidates(
         )
 
 
+# `!secret key`, quoted or not, also inside flow collections. A text scan, not a
+# YAML parse, so it can also match comments; that errs toward more secrets.
+_SECRET_REFERENCE_RE = re.compile(r"""!secret\s+['"]?([^\s'",\]}]+)""")
+
+
+def find_secret_references(text: str) -> set[str]:
+    """Return the ``!secret <key>`` names referenced in a YAML document text."""
+    return {match.group(1) for match in _SECRET_REFERENCE_RE.finditer(text)}
+
+
 def force_load_include_files(
     obj: Any,
     *,
@@ -858,6 +875,7 @@ class ESPHomeLoaderMixin:
             )
         val = secrets[node.value]
         _SECRET_VALUES[str(val)] = node.value
+        _LOADED_SECRET_NAMES.add(node.value)
         return val
 
     @_add_data_ref
@@ -1030,6 +1048,7 @@ def load_yaml(
     if clear_secrets:
         _SECRET_VALUES.clear()
         _SECRET_CACHE.clear()
+        _LOADED_SECRET_NAMES.clear()
     return _load_yaml_internal(fname, track_document_range=track_document_range)
 
 
@@ -1135,6 +1154,38 @@ def _load_yaml_internal_with_type(
         loader.dispose()
 
 
+def registered_secret_names() -> set[str]:
+    """Names of all ``!secret`` keys the loader has seen since the last clear."""
+    return set(_LOADED_SECRET_NAMES)
+
+
+@contextmanager
+def secret_values_registered(values: dict[str, str]) -> Generator[set[str]]:
+    """Make :func:`dump` write these values as ``!secret <name>`` while active.
+
+    Real secrets already registered win. Yields the set of ``!secret`` names the
+    dumper emits. A wrapped value (``!extend``, ``!remove``, ``!include``) loses its
+    tag rather than leaking.
+    """
+    global _EMITTED_SECRET_NAMES  # noqa: PLW0603
+    added = {v: n for v, n in values.items() if v not in _SECRET_VALUES}
+    _SECRET_VALUES.update(added)
+    outer = _EMITTED_SECRET_NAMES
+    emitted: set[str] = set()
+    _EMITTED_SECRET_NAMES = emitted
+    try:
+        yield emitted
+    finally:
+        _EMITTED_SECRET_NAMES = outer
+        if outer is not None:
+            outer |= emitted
+        # Only drop mappings this context still owns; a real !secret loaded
+        # meanwhile may have registered the same value under its own name.
+        for value, name in added.items():
+            if _SECRET_VALUES.get(value) == name:
+                del _SECRET_VALUES[value]
+
+
 def dump(
     dict_,
     show_secrets=False,
@@ -1152,6 +1203,7 @@ def dump(
     if show_secrets:
         _SECRET_VALUES.clear()
         _SECRET_CACHE.clear()
+        _LOADED_SECRET_NAMES.clear()
 
     # Per-call subclass so the flags don't leak across calls.
     # (``_SECRET_VALUES`` / ``_SECRET_CACHE`` remain module globals; YAML
@@ -1349,7 +1401,10 @@ class ESPHomeDumper(yaml.SafeDumper):
         return node
 
     def represent_secret(self, value):
-        return self.represent_scalar(tag="!secret", value=_SECRET_VALUES[str(value)])
+        name = _SECRET_VALUES[str(value)]
+        if _EMITTED_SECRET_NAMES is not None:
+            _EMITTED_SECRET_NAMES.add(name)
+        return self.represent_scalar(tag="!secret", value=name)
 
     def represent_stringify(self, value):
         if is_secret(value):
@@ -1433,14 +1488,33 @@ class ESPHomeDumper(yaml.SafeDumper):
             return self.represent_secret(value.value)
         return self.represent_scalar(tag="!lambda", value=value.value, style="|")
 
+    def _represent_tagged_secret(self, tag: str, value: Any) -> yaml.ScalarNode:
+        # Writing the secret beats leaking the value, but the tag is lost, so
+        # the dumped config no longer means the same thing when loaded back.
+        _LOGGER.warning(
+            "%s value matches secret '%s'; written as !secret, dropping the %s tag",
+            tag,
+            _SECRET_VALUES[str(value)],
+            tag,
+        )
+        return self.represent_secret(value)
+
     def represent_extend(self, value):
+        # Consult is_secret like the other scalar representers so a payload
+        # equal to a registered secret is never written out in cleartext.
+        if is_secret(value.value):
+            return self._represent_tagged_secret("!extend", value.value)
         return self.represent_scalar(tag="!extend", value=value.value)
 
     def represent_remove(self, value):
+        if is_secret(value.value):
+            return self._represent_tagged_secret("!remove", value.value)
         return self.represent_scalar(tag="!remove", value=value.value)
 
     def represent_include_file(self, value):
         if value.vars or value.condition is not None:
+            # The mapping values route through the regular representers,
+            # which already consult is_secret.
             mapping = {"file": value.file}
             if value.vars:
                 mapping["vars"] = value.vars
@@ -1449,6 +1523,8 @@ class ESPHomeDumper(yaml.SafeDumper):
             return self.represent_mapping(
                 tag="!include", mapping=mapping, flow_style=False
             )
+        if is_secret(value.file):
+            return self._represent_tagged_secret("!include", value.file)
         return self.represent_scalar(tag="!include", value=value.file)
 
     def represent_id(self, value):

tests/unit_tests/conftest.py

@@ -17,6 +17,7 @@ from unittest.mock import Mock, patch
 
 import pytest
 
+from esphome import yaml_util
 from esphome.core import CORE
 from esphome.espidf import toolchain
 
@@ -49,6 +50,18 @@ def reset_core():
     CORE.reset()
 
 
+@pytest.fixture(autouse=True)
+def clear_yaml_secrets() -> Generator[None]:
+    """Isolate the yaml_util secrets registry between tests."""
+    yaml_util._SECRET_VALUES.clear()
+    yaml_util._SECRET_CACHE.clear()
+    yaml_util._LOADED_SECRET_NAMES.clear()
+    yield
+    yaml_util._SECRET_VALUES.clear()
+    yaml_util._SECRET_CACHE.clear()
+    yaml_util._LOADED_SECRET_NAMES.clear()
+
+
 @pytest.fixture
 def fixture_path() -> Path:
     """

tests/unit_tests/test_bundle.py

@@ -139,10 +139,17 @@ def test_find_used_secret_keys_no_secrets(tmp_path: Path) -> None:
     assert keys == set()
 
 
-def test_find_used_secret_keys_missing_file(tmp_path: Path) -> None:
+def test_find_used_secret_keys_missing_file(
+    tmp_path: Path, caplog: pytest.LogCaptureFixture
+) -> None:
     missing = tmp_path / "does_not_exist.yaml"
-    keys = _find_used_secret_keys([missing])
+    with caplog.at_level("WARNING", logger="esphome.bundle"):
+        keys = _find_used_secret_keys([missing])
     assert keys == set()
+    assert any(
+        "Could not scan" in r.message and "does_not_exist.yaml" in r.message
+        for r in caplog.records
+    )
 
 
 def test_find_used_secret_keys_deduplicates(tmp_path: Path) -> None:

tests/unit_tests/test_yaml_util.py

@@ -26,16 +26,6 @@ from esphome.yaml_util import (
 )
 
 
-@pytest.fixture(autouse=True)
-def clear_secrets_cache() -> None:
-    """Clear the secrets cache before each test."""
-    yaml_util._SECRET_VALUES.clear()
-    yaml_util._SECRET_CACHE.clear()
-    yield
-    yaml_util._SECRET_VALUES.clear()
-    yaml_util._SECRET_CACHE.clear()
-
-
 @pytest.fixture(autouse=True)
 def clear_core_frontmatter() -> None:
     """Reset CORE.frontmatter between tests."""
@@ -1973,6 +1963,72 @@ def test_dump__redaction_flag_does_not_leak_between_calls() -> None:
     assert "\\033[8m" in redacted_again
 
 
+def test_secret_values_registered_swaps_scalars_in_dump() -> None:
+    """Registered value→name mappings make dump() emit `!secret <name>` for
+    matching scalars, and are removed again on exit."""
+    with yaml_util.secret_values_registered({"hunter2": "wifi_password"}):
+        out = yaml_util.dump({"password": make_data_base("hunter2")})
+        assert "password: !secret 'wifi_password'" in out
+        assert "hunter2" not in out
+    out_after = yaml_util.dump({"password": make_data_base("hunter2")})
+    assert "hunter2" in out_after
+    assert "!secret" not in out_after
+    assert yaml_util.is_secret("hunter2") is None
+
+
+def test_secret_values_registered_collects_emitted_names() -> None:
+    """The context yields a set recording exactly the `!secret` names the
+    dumper emitted, including real secrets, and not names never swapped."""
+    yaml_util._SECRET_VALUES["real_value"] = "real_name"
+    with yaml_util.secret_values_registered({"hunter2": "wifi_password"}) as emitted:
+        yaml_util.dump(
+            {
+                "password": make_data_base("hunter2"),
+                "key": make_data_base("real_value"),
+                "plain": make_data_base("nothing"),
+            }
+        )
+    assert emitted == {"wifi_password", "real_name"}
+
+
+def test_secret_values_registered_does_not_clobber_real_secrets() -> None:
+    """A value already mapped by a real `!secret` keeps its original name."""
+    yaml_util._SECRET_VALUES["hunter2"] = "original_name"
+    with yaml_util.secret_values_registered({"hunter2": "generated_name"}):
+        out = yaml_util.dump({"password": make_data_base("hunter2")})
+        assert "!secret 'original_name'" in out
+    # The pre-existing mapping survives the context exit.
+    assert yaml_util.is_secret("hunter2") == "original_name"
+
+
+def test_secret_values_registered_keeps_a_real_secret_loaded_meanwhile() -> None:
+    """A real `!secret` registered inside the context survives its exit."""
+    with yaml_util.secret_values_registered({"hunter2": "generated_name"}):
+        yaml_util._SECRET_VALUES["hunter2"] = "real_name"
+    assert yaml_util.is_secret("hunter2") == "real_name"
+
+
+def test_secret_values_registered_nested_keeps_the_outer_collector() -> None:
+    """The outer context also collects names emitted inside an inner one."""
+    with yaml_util.secret_values_registered({"outer_value": "outer"}) as outer:
+        with yaml_util.secret_values_registered({"inner_value": "inner"}) as inner:
+            yaml_util.dump({"a": make_data_base("inner_value")})
+        yaml_util.dump({"b": make_data_base("outer_value")})
+    assert inner == {"inner"}
+    assert outer == {"inner", "outer"}
+
+
+def test_registered_secret_names(tmp_path: Path) -> None:
+    """Every loaded name is returned, even two sharing a value, but not names
+    registered only for a dump."""
+    (tmp_path / "secrets.yaml").write_text("wifi_password: same\nap_password: same\n")
+    main = tmp_path / "main.yaml"
+    main.write_text("a: !secret wifi_password\nb: !secret ap_password\n")
+    yaml_util.load_yaml(main)
+    with yaml_util.secret_values_registered({"other": "temp_name"}):
+        assert yaml_util.registered_secret_names() == {"wifi_password", "ap_password"}
+
+
 @pytest.fixture(autouse=True)
 def clear_dropped_merge_keys() -> None:
     """Reset the dropped-merge-key queue between tests."""
@@ -2018,6 +2074,42 @@ def test_merge_include_no_overlap_records_nothing(tmp_path: Path) -> None:
     assert yaml_util.take_dropped_merge_keys() == []
 
 
+def test_wrapper_representers_consult_is_secret(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """!extend / !remove payloads and scalar !include paths equal to a
+    registered secret are swapped, never written in cleartext, and the lost
+    tag is logged."""
+    from esphome.config_helpers import Extend, Remove
+
+    with (
+        caplog.at_level("WARNING", logger="esphome.yaml_util"),
+        yaml_util.secret_values_registered({"hunter2": "the_secret"}),
+    ):
+        out = yaml_util.dump(
+            {
+                "a": Extend("hunter2"),
+                "b": Remove("hunter2"),
+                "c": Extend("plain_id"),
+            }
+        )
+    assert out.count("!secret 'the_secret'") == 2
+    assert "hunter2" not in out
+    assert "!extend 'plain_id'" in out
+    lost = [r.message for r in caplog.records if "dropping the" in r.message]
+    assert len(lost) == 2
+    assert all("hunter2" not in m for m in lost)
+
+
+def test_scalar_include_path_equal_to_secret_is_swapped() -> None:
+    """A scalar !include whose path equals a registered secret is swapped,
+    matching the other wrapper representers."""
+    include = yaml_util.IncludeFile(Path("/fake/main.yaml"), "hunter2", lambda _: {})
+    with yaml_util.secret_values_registered({"hunter2": "the_secret"}):
+        out = yaml_util.dump({"key": include})
+    assert out == "key: !secret 'the_secret'\n"
+
+
 # ---------------------------------------------------------------------------
 # track_document_range=False (validated-config-cache fast path)
 # ---------------------------------------------------------------------------
@@ -2135,3 +2227,9 @@ def test_load_yaml_fast_mode_records_dropped_merge_keys(
 
     yaml_util.load_yaml(yaml_file, track_document_range=False)
     assert yaml_util.take_dropped_merge_keys() == [("port", str(yaml_file))]
+
+
+def test_find_secret_references_in_flow_collections() -> None:
+    """Flow mapping and sequence punctuation is not part of the secret name."""
+    text = "a: {password: !secret wifi_pw}\nb: [!secret one, !secret 'two']\n"
+    assert yaml_util.find_secret_references(text) == {"wifi_pw", "one", "two"}