[ota] Offer encryption with the api key so enabling it works over OTA (#18979)

tests/integration/conftest.py

@@ -162,6 +162,13 @@ def integration_test_dir() -> Generator[Path]:
         yield Path(tmpdir)
 
 
+@pytest.fixture
+def isolated_preferences(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    """Host preferences persist per device name; give the test its own so a
+    provisioned key never leaks into another run."""
+    monkeypatch.setenv("ESPHOME_PREFDIR", str(tmp_path / "prefs"))
+
+
 @pytest.fixture
 def reserved_tcp_port() -> Generator[tuple[int, socket.socket]]:
     """Reserve an unused TCP port by holding the socket open."""

tests/integration/const.py

@@ -9,6 +9,13 @@ API_CONNECTION_TIMEOUT = 30.0  # seconds
 PORT_WAIT_TIMEOUT = 30.0  # seconds
 PORT_POLL_INTERVAL = 0.1  # seconds
 
+# The well-known all-zeros provisioning PSK, a key to provision over it, and
+# the time the device takes to activate a newly saved key (100 ms timer plus
+# margin)
+ZERO_PSK = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA="
+PROVISIONING_PSK = b"bm5ubm5ubm5ubm5ubm5ubm5ubm5ubm5ubm5ubm5ubm4="
+KEY_ACTIVATION_DELAY = 0.5  # seconds
+
 # Process shutdown timeouts
 SIGINT_TIMEOUT = 5.0  # seconds
 SIGTERM_TIMEOUT = 2.0  # seconds

tests/integration/fixtures/host_ota_api_key_offer_with_password.yaml

@@ -0,0 +1,12 @@
+esphome:
+  name: host-ota-test
+host:
+api:
+  encryption:
+    key: "AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8="
+ota:
+  - platform: esphome
+    port: __OTA_PORT__
+    password: "hunter2"
+logger:
+  level: DEBUG

tests/integration/fixtures/host_ota_provisioned_api_key.yaml

@@ -0,0 +1,10 @@
+esphome:
+  name: host-ota-test
+host:
+api:
+  encryption:
+ota:
+  - platform: esphome
+    port: __OTA_PORT__
+logger:
+  level: DEBUG

tests/integration/test_api_zero_psk_provisioning.py

@@ -10,34 +10,40 @@ from __future__ import annotations
 
 import asyncio
 import base64
+import socket
 
 from aioesphomeapi import InvalidEncryptionKeyAPIError, RequiresEncryptionAPIError
 import pytest
 
-from .types import APIClientConnectedFactory, RunCompiledFunction
+from .conftest import run_binary_and_wait_for_port
+from .const import KEY_ACTIVATION_DELAY, LOCALHOST, PROVISIONING_PSK, ZERO_PSK
+from .types import (
+    APIClientConnectedFactory,
+    CompileFunction,
+    ConfigWriter,
+    RunCompiledFunction,
+)
 
-# The well-known provisioning PSK: base64 of 32 zero bytes
-ZERO_PSK = base64.b64encode(bytes(32)).decode()
-# A real key to provision
-NEW_KEY = base64.b64encode(b"n" * 32)
-# Time for the device to activate a newly saved key (100ms timer plus margin)
-KEY_ACTIVATION_DELAY = 0.5
-
-
-@pytest.fixture(autouse=True)
-def isolated_preferences(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None:
-    """Keep host preferences per-test so every run starts unprovisioned."""
-    monkeypatch.setenv("ESPHOME_PREFDIR", str(tmp_path / "prefs"))
+pytestmark = pytest.mark.usefixtures("isolated_preferences")
+NEW_KEY = PROVISIONING_PSK
 
 
 @pytest.mark.asyncio
 async def test_api_zero_psk_provisioning(
     yaml_config: str,
-    run_compiled: RunCompiledFunction,
+    write_yaml_config: ConfigWriter,
+    compile_esphome: CompileFunction,
+    reserved_tcp_port: tuple[int, socket.socket],
     api_client_connected: APIClientConnectedFactory,
 ) -> None:
-    """Exercise the reject paths, then provision a key over the zero-PSK channel."""
-    async with run_compiled(yaml_config):
+    """Exercise the reject paths, provision a key over the zero-PSK channel,
+    and check the key comes back from preferences on the next boot."""
+    port, port_socket = reserved_tcp_port
+    config_path = await write_yaml_config(yaml_config)
+    binary_path = await compile_esphome(config_path)
+    port_socket.close()
+
+    async with run_binary_and_wait_for_port(binary_path, LOCALHOST, port):
         # --- Pre-provisioning reject paths (device state is unchanged) ---
 
         # A wrong (non-zero) PSK fails against the zero provisioning PSK
@@ -97,6 +103,19 @@ async def test_api_zero_psk_provisioning(
             async with api_client_connected(timeout=5) as client:
                 await client.device_info()
 
+    # The key is loaded from preferences on the next boot
+    lines: list[str] = []
+    async with run_binary_and_wait_for_port(
+        binary_path, LOCALHOST, port, line_callback=lines.append
+    ):
+        async with api_client_connected(noise_psk=NEW_KEY.decode()) as client:
+            device_info = await client.device_info()
+            assert device_info.api_encryption_provisionable is False
+        with pytest.raises(InvalidEncryptionKeyAPIError):
+            async with api_client_connected(noise_psk=ZERO_PSK, timeout=5) as client:
+                await client.device_info()
+    assert any("Loaded saved Noise PSK" in line for line in lines)
+
 
 @pytest.mark.asyncio
 async def test_api_zero_psk_provisioning_plaintext(

tests/integration/test_host_ota.py

@@ -8,9 +8,12 @@ instance covers the FD_CLOEXEC path.
 from __future__ import annotations
 
 import asyncio
+import base64
 from collections.abc import Generator
 from contextlib import contextmanager
+from dataclasses import dataclass
 import functools
+from pathlib import Path
 import socket
 
 import pytest
@@ -18,10 +21,18 @@ import pytest
 from esphome import espota2
 
 from .conftest import run_binary, wait_and_connect_api_client
-from .const import LOCALHOST, PORT_POLL_INTERVAL, PORT_WAIT_TIMEOUT
-from .types import CompileFunction, ConfigWriter
+from .const import (
+    KEY_ACTIVATION_DELAY,
+    LOCALHOST,
+    PORT_POLL_INTERVAL,
+    PORT_WAIT_TIMEOUT,
+    PROVISIONING_PSK,
+    ZERO_PSK,
+)
+from .types import APIClientConnectedFactory, CompileFunction, ConfigWriter
 
 DEVICE_NAME = "host-ota-test"
+API_KEY = "AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8="
 
 
 @contextmanager
@@ -35,6 +46,14 @@ def _reserve_port() -> Generator[tuple[int, socket.socket]]:
         s.close()
 
 
+async def _wait_for_line(lines: list[str], needle: str, timeout: float = 5.0) -> None:
+    """The config dump prints after every setup, a little after the api port
+    opens, so wait for it rather than assert on the lines seen so far."""
+    async with asyncio.timeout(timeout):
+        while not any(needle in line for line in lines):
+            await asyncio.sleep(PORT_POLL_INTERVAL)
+
+
 async def _wait_for_port(host: str, port: int, timeout: float) -> None:
     """Poll until a TCP port accepts connections, or raise TimeoutError."""
     loop = asyncio.get_running_loop()
@@ -51,6 +70,102 @@ async def _wait_for_port(host: str, port: int, timeout: float) -> None:
     raise TimeoutError(f"Port {port} on {host} did not open within {timeout}s")
 
 
+async def _build(
+    yaml_config: str,
+    write_yaml_config: ConfigWriter,
+    compile_esphome: CompileFunction,
+    reserved_tcp_port: tuple[int, socket.socket],
+) -> tuple[int, int, Path]:
+    """Reserve an OTA port, compile the fixture with it, and release both
+    ports right before the binary is started."""
+    api_port, api_socket = reserved_tcp_port
+    with _reserve_port() as (ota_port, ota_socket):
+        config_path = await write_yaml_config(
+            yaml_config.replace("__OTA_PORT__", str(ota_port))
+        )
+        binary_path = await compile_esphome(config_path)
+        api_socket.close()
+        ota_socket.close()
+    return api_port, ota_port, binary_path
+
+
+async def _run_ota(
+    ota_port: int,
+    password: str | None,
+    binary_path: Path,
+    noise_psk: str | None,
+    plaintext_fallback: bool = False,
+) -> int:
+    """espota2 is blocking; run it in the executor and return its exit code."""
+    rc, _ = await asyncio.get_running_loop().run_in_executor(
+        None,
+        functools.partial(
+            espota2.run_ota,
+            LOCALHOST,
+            ota_port,
+            password,
+            binary_path,
+            noise_psk=noise_psk,
+            plaintext_fallback=plaintext_fallback,
+        ),
+    )
+    return rc
+
+
+@dataclass
+class _Device:
+    """A running host binary and the checks every successful OTA repeats:
+    a safe reboot, the api port back up, and the pid preserved by execv."""
+
+    api_port: int
+    ota_port: int
+    binary_path: Path
+    proc: asyncio.subprocess.Process | None = None
+    reboots: int = 0
+
+    def __post_init__(self) -> None:
+        self._rebooted = asyncio.Event()
+
+    def on_log(self, line: str) -> None:
+        if "Rebooting safely" in line:
+            self.reboots += 1
+            self._rebooted.set()
+
+    async def wait_reboot(self, count: int, timeout: float = 10.0) -> None:
+        async with asyncio.timeout(timeout):
+            while self.reboots < count:
+                self._rebooted.clear()
+                await self._rebooted.wait()
+
+    async def ota(
+        self,
+        password: str | None,
+        noise_psk: str | None,
+        msg: str,
+        plaintext_fallback: bool = False,
+    ) -> None:
+        """Upload, then expect the re-exec with the pid preserved."""
+        pid_before = self.proc.pid
+        expected_reboots = self.reboots + 1
+        rc = await _run_ota(
+            self.ota_port, password, self.binary_path, noise_psk, plaintext_fallback
+        )
+        assert rc == 0, msg
+        await self.wait_reboot(expected_reboots)
+        await _wait_for_port(LOCALHOST, self.api_port, PORT_WAIT_TIMEOUT)
+        assert self.proc.returncode is None, "process exited instead of execing"
+        assert self.proc.pid == pid_before
+
+    async def refused_ota(
+        self, password: str | None, noise_psk: str | None, msg: str
+    ) -> None:
+        """Upload must fail and the device must keep running."""
+        rc = await _run_ota(self.ota_port, password, self.binary_path, noise_psk)
+        assert rc == 1, msg
+        await asyncio.sleep(0.5)
+        assert self.proc.returncode is None, "process died on rejected OTA"
+
+
 @pytest.mark.asyncio
 async def test_host_ota_self_update(
     yaml_config: str,
@@ -59,57 +174,34 @@ async def test_host_ota_self_update(
     reserved_tcp_port: tuple[int, socket.socket],
 ) -> None:
     """Self-OTA: upload the running binary back to itself, expect re-exec."""
-    api_port, api_socket = reserved_tcp_port
-    with _reserve_port() as (ota_port, ota_socket):
-        yaml_config = yaml_config.replace("__OTA_PORT__", str(ota_port))
-        config_path = await write_yaml_config(yaml_config)
-        binary_path = await compile_esphome(config_path)
-        api_socket.close()
-        ota_socket.close()
+    dev = _Device(
+        *await _build(
+            yaml_config, write_yaml_config, compile_esphome, reserved_tcp_port
+        )
+    )
+    staged = asyncio.Event()
 
-        loop = asyncio.get_running_loop()
-        ota_staged = loop.create_future()
-        rebooted = loop.create_future()
+    def on_log(line: str) -> None:
+        if "OTA staged at" in line:
+            staged.set()
+        dev.on_log(line)
 
-        def on_log(line: str) -> None:
-            if not ota_staged.done() and "OTA staged at" in line:
-                ota_staged.set_result(True)
-            if not rebooted.done() and "Rebooting safely" in line:
-                rebooted.set_result(True)
+    async with run_binary(dev.binary_path, line_callback=on_log) as (proc, _lines):
+        dev.proc = proc
+        await _wait_for_port(LOCALHOST, dev.api_port, PORT_WAIT_TIMEOUT)
+        async with wait_and_connect_api_client(port=dev.api_port) as client:
+            info_before = await client.device_info()
+            assert info_before.name == DEVICE_NAME
 
-        async with run_binary(binary_path, line_callback=on_log) as (proc, _lines):
-            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
-            pid_before = proc.pid
-            async with wait_and_connect_api_client(port=api_port) as client:
-                info_before = await client.device_info()
-                assert info_before.name == DEVICE_NAME
+        await dev.ota(None, None, "espota2 reported failure")
+        assert staged.is_set()
 
-            # espota2 is blocking; run in executor.
-            rc, _ = await loop.run_in_executor(
-                None, espota2.run_ota, LOCALHOST, ota_port, None, binary_path
-            )
-            assert rc == 0, "espota2 reported failure"
+        async with wait_and_connect_api_client(port=dev.api_port) as client:
+            info_after = await client.device_info()
+            assert info_after.name == info_before.name
 
-            await asyncio.wait_for(ota_staged, timeout=10.0)
-            await asyncio.wait_for(rebooted, timeout=10.0)
-            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
-
-            # execv preserves pid; mismatch means external respawn.
-            assert proc.returncode is None, "process exited instead of execing"
-            assert proc.pid == pid_before
-
-            async with wait_and_connect_api_client(port=api_port) as client:
-                info_after = await client.device_info()
-                assert info_after.name == DEVICE_NAME
-                assert info_after.name == info_before.name
-
-            # Second OTA: catches FD_CLOEXEC regressions (EADDRINUSE on rebind).
-            rc, _ = await loop.run_in_executor(
-                None, espota2.run_ota, LOCALHOST, ota_port, None, binary_path
-            )
-            assert rc == 0, "second OTA failed -- listener leaked across execv"
-            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
-            assert proc.pid == pid_before
+        # Second OTA: catches FD_CLOEXEC regressions (EADDRINUSE on rebind).
+        await dev.ota(None, None, "second OTA failed -- listener leaked across execv")
 
 
 @pytest.mark.asyncio
@@ -121,51 +213,110 @@ async def test_host_ota_encrypted(
 ) -> None:
     """Encrypted self-OTA succeeds; a plaintext upload to the same device fails."""
     pytest.importorskip("aioesphomeapi.noise")
-    noise_psk = "AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8="
-    api_port, api_socket = reserved_tcp_port
-    with _reserve_port() as (ota_port, ota_socket):
-        yaml_config = yaml_config.replace("__OTA_PORT__", str(ota_port))
-        config_path = await write_yaml_config(yaml_config)
-        binary_path = await compile_esphome(config_path)
-        api_socket.close()
-        ota_socket.close()
+    dev = _Device(
+        *await _build(
+            yaml_config, write_yaml_config, compile_esphome, reserved_tcp_port
+        )
+    )
+    async with run_binary(dev.binary_path, line_callback=dev.on_log) as (proc, _lines):
+        dev.proc = proc
+        await _wait_for_port(LOCALHOST, dev.api_port, PORT_WAIT_TIMEOUT)
+        await dev.refused_ota(
+            None, None, "plaintext upload to an encrypted device must fail"
+        )
+        await dev.ota(None, API_KEY, "encrypted OTA reported failure")
 
-        loop = asyncio.get_running_loop()
-        rebooted = loop.create_future()
 
-        def on_log(line: str) -> None:
-            if not rebooted.done() and "Rebooting safely" in line:
-                rebooted.set_result(True)
+@pytest.mark.asyncio
+async def test_host_ota_api_key_offer_with_password(
+    yaml_config: str,
+    write_yaml_config: ConfigWriter,
+    compile_esphome: CompileFunction,
+    reserved_tcp_port: tuple[int, socket.socket],
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """With only an api key the device offers encryption without requiring
+    it: the password still guards plaintext uploads, the key alone
+    authenticates an encrypted one, and until 2027.3.0 a failed encrypted
+    attempt falls back to plaintext."""
+    pytest.importorskip("aioesphomeapi.noise")
+    wrong_key = base64.b64encode(b"w" * 32).decode()
+    dev = _Device(
+        *await _build(
+            yaml_config, write_yaml_config, compile_esphome, reserved_tcp_port
+        )
+    )
+    async with run_binary(dev.binary_path, line_callback=dev.on_log) as (proc, lines):
+        dev.proc = proc
+        await _wait_for_port(LOCALHOST, dev.api_port, PORT_WAIT_TIMEOUT)
+        await _wait_for_line(lines, "Encryption: offered")
 
-        async with run_binary(binary_path, line_callback=on_log) as (proc, _lines):
-            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
-            pid_before = proc.pid
+        await dev.refused_ota(
+            None, None, "plaintext upload without the password must fail"
+        )
+        await dev.ota(
+            "hunter2", None, "plaintext upload with the password must succeed"
+        )
+        await dev.ota(None, API_KEY, "encrypted upload with the api key must succeed")
 
-            # A plaintext upload must be refused with the device unharmed
-            rc, _ = await loop.run_in_executor(
-                None, espota2.run_ota, LOCALHOST, ota_port, None, binary_path
+        # Remove before 2027.3.0: a wrong key falls back to plaintext, which
+        # the password still guards
+        with caplog.at_level("WARNING", logger="esphome.espota2"):
+            await dev.ota(
+                "hunter2",
+                wrong_key,
+                "the plaintext retry with the password must succeed",
+                plaintext_fallback=True,
             )
-            assert rc == 1, "plaintext upload to an encrypted device must fail"
-            await asyncio.sleep(0.5)
-            assert proc.returncode is None, "process died on rejected plaintext OTA"
+        assert any("Retrying in plaintext" in r.message for r in caplog.records)
+        await dev.ota(
+            None,
+            API_KEY,
+            "the right api key encrypts without touching the fallback",
+            plaintext_fallback=True,
+        )
 
-            # The encrypted upload goes through and the device re-execs
-            rc, _ = await loop.run_in_executor(
-                None,
-                functools.partial(
-                    espota2.run_ota,
-                    LOCALHOST,
-                    ota_port,
-                    None,
-                    binary_path,
-                    noise_psk=noise_psk,
-                ),
-            )
-            assert rc == 0, "encrypted OTA reported failure"
-            await asyncio.wait_for(rebooted, timeout=10.0)
-            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
-            assert proc.returncode is None, "process exited instead of execing"
-            assert proc.pid == pid_before
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("isolated_preferences")
+async def test_host_ota_provisioned_api_key(
+    yaml_config: str,
+    write_yaml_config: ConfigWriter,
+    compile_esphome: CompileFunction,
+    reserved_tcp_port: tuple[int, socket.socket],
+    api_client_connected: APIClientConnectedFactory,
+) -> None:
+    """A key provisioned over the api feeds the OTA offer: plaintext works
+    while unprovisioned, the provisioned key encrypts, the key loaded from
+    preferences on the next boot keeps encrypting, and plaintext stays
+    accepted because only the ota block requires encryption."""
+    pytest.importorskip("aioesphomeapi.noise")
+    dev = _Device(
+        *await _build(
+            yaml_config, write_yaml_config, compile_esphome, reserved_tcp_port
+        )
+    )
+    async with run_binary(dev.binary_path, line_callback=dev.on_log) as (proc, lines):
+        dev.proc = proc
+        await _wait_for_port(LOCALHOST, dev.api_port, PORT_WAIT_TIMEOUT)
+        await _wait_for_line(lines, "once the api key is provisioned")
+
+        await dev.ota(
+            None, None, "plaintext upload to an unprovisioned device must succeed"
+        )
+
+        async with api_client_connected(
+            port=dev.api_port, noise_psk=ZERO_PSK
+        ) as client:
+            assert await client.noise_encryption_set_key(PROVISIONING_PSK) is True
+        await asyncio.sleep(KEY_ACTIVATION_DELAY)
+
+        key = PROVISIONING_PSK.decode()
+        await dev.ota(
+            None, key, "encrypted upload with the provisioned key must succeed"
+        )
+        await dev.ota(None, key, "the key loaded at boot must feed the OTA offer")
+        await dev.ota(None, None, "plaintext must stay accepted on an offering device")
 
 
 @pytest.mark.asyncio
@@ -177,33 +328,25 @@ async def test_host_ota_rejects_garbage(
     integration_test_dir,
 ) -> None:
     """Bogus payload is rejected and the device keeps running."""
-    api_port, api_socket = reserved_tcp_port
-    with _reserve_port() as (ota_port, ota_socket):
-        yaml_config = yaml_config.replace("__OTA_PORT__", str(ota_port))
-        config_path = await write_yaml_config(yaml_config)
-        binary_path = await compile_esphome(config_path)
+    dev = _Device(
+        *await _build(
+            yaml_config, write_yaml_config, compile_esphome, reserved_tcp_port
+        )
+    )
+    # 192 bytes that are neither ELF nor Mach-O.
+    bogus_path = integration_test_dir / "bogus.bin"
+    bogus_path.write_bytes(b"NOT-AN-EXECUTABLE-AT-ALL" * 8)
 
-        # 192 bytes that are neither ELF nor Mach-O.
-        bogus_path = integration_test_dir / "bogus.bin"
-        bogus_path.write_bytes(b"NOT-AN-EXECUTABLE-AT-ALL" * 8)
+    async with run_binary(dev.binary_path) as (proc, _lines):
+        dev.proc = proc
+        await _wait_for_port(LOCALHOST, dev.api_port, PORT_WAIT_TIMEOUT)
+        pid_before = proc.pid
+        rc = await _run_ota(dev.ota_port, None, bogus_path, None)
+        assert rc == 1
+        await asyncio.sleep(0.5)
+        assert proc.returncode is None, "process died on rejected OTA"
+        assert proc.pid == pid_before
 
-        api_socket.close()
-        ota_socket.close()
-
-        async with run_binary(binary_path) as (proc, _lines):
-            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
-            pid_before = proc.pid
-
-            loop = asyncio.get_running_loop()
-            rc, _ = await loop.run_in_executor(
-                None, espota2.run_ota, LOCALHOST, ota_port, None, bogus_path
-            )
-            assert rc == 1
-
-            await asyncio.sleep(0.5)
-            assert proc.returncode is None, "process died on rejected OTA"
-            assert proc.pid == pid_before
-
-            async with wait_and_connect_api_client(port=api_port) as client:
-                info = await client.device_info()
-                assert info.name == DEVICE_NAME
+        async with wait_and_connect_api_client(port=dev.api_port) as client:
+            info = await client.device_info()
+            assert info.name == DEVICE_NAME