esphome/tests/integration/test_api_zero_psk_provisioning.py

"""Integration tests for provisioning the encryption key over a zero-PSK connection.

A device with `api: encryption:` but no key accepts Noise handshakes using the
well-known all-zeros PSK. The ephemeral X25519 exchange protects the key from
passive sniffing while it is provisioned; plaintext provisioning still works
but is deprecated.
"""

from __future__ import annotations

import asyncio
import base64
import socket

from aioesphomeapi import InvalidEncryptionKeyAPIError, RequiresEncryptionAPIError
import pytest

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,
)

pytestmark = pytest.mark.usefixtures("isolated_preferences")
NEW_KEY = PROVISIONING_PSK


@pytest.mark.asyncio
async def test_api_zero_psk_provisioning(
    yaml_config: str,
    write_yaml_config: ConfigWriter,
    compile_esphome: CompileFunction,
    reserved_tcp_port: tuple[int, socket.socket],
    api_client_connected: APIClientConnectedFactory,
) -> None:
    """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
        with pytest.raises(InvalidEncryptionKeyAPIError):
            async with api_client_connected(
                noise_psk=base64.b64encode(b"w" * 32).decode(), timeout=5
            ) as client:
                await client.device_info()

        # A plaintext client and a zero-PSK client can be connected at the
        # same time while the device is unprovisioned
        async with (
            api_client_connected() as plaintext_client,
            api_client_connected(noise_psk=ZERO_PSK) as noise_client,
        ):
            plaintext_info = await plaintext_client.device_info()
            noise_info = await noise_client.device_info()
            # Both transports advertise provisioning support so old and new
            # clients can decide how to provision
            assert plaintext_info.api_encryption_provisionable is True
            assert noise_info.api_encryption_provisionable is True

            # The all-zeros key is reserved as the provisioning PSK and is
            # rejected on both transports
            zero_key = base64.b64encode(bytes(32))
            assert await noise_client.noise_encryption_set_key(zero_key) is False
            assert await plaintext_client.noise_encryption_set_key(zero_key) is False

        # --- Provision over the zero-PSK channel ---

        # The unprovisioned device accepts the all-zeros PSK; the handshake's
        # ephemeral-ephemeral DH encrypts everything that follows
        async with api_client_connected(noise_psk=ZERO_PSK) as client:
            device_info = await client.device_info()
            assert device_info.name == "zero-psk-provision-test"
            assert device_info.api_encryption_supported is True
            assert device_info.api_encryption_provisionable is True

            assert await client.noise_encryption_set_key(NEW_KEY) is True

        # The device activates the new key shortly after responding
        await asyncio.sleep(KEY_ACTIVATION_DELAY)

        # The new key now works, and the device is no longer provisionable
        async with api_client_connected(noise_psk=NEW_KEY.decode()) as client:
            device_info = await client.device_info()
            assert device_info.name == "zero-psk-provision-test"
            assert device_info.api_encryption_provisionable is False

        # The zero PSK no longer works
        with pytest.raises(InvalidEncryptionKeyAPIError):
            async with api_client_connected(noise_psk=ZERO_PSK, timeout=5) as client:
                await client.device_info()

        # Plaintext no longer works
        with pytest.raises(RequiresEncryptionAPIError):
            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(
    yaml_config: str,
    run_compiled: RunCompiledFunction,
    api_client_connected: APIClientConnectedFactory,
) -> None:
    """The legacy plaintext provisioning path still works and warns."""
    log_lines: list[str] = []
    async with run_compiled(yaml_config, line_callback=log_lines.append):
        async with api_client_connected() as client:
            device_info = await client.device_info()
            assert device_info.name == "zero-psk-plaintext-test"

            assert await client.noise_encryption_set_key(NEW_KEY) is True

        await asyncio.sleep(KEY_ACTIVATION_DELAY)

        # The deprecation warning was logged
        assert any("deprecated" in line for line in log_lines)

        # The new key works; the zero PSK does not
        async with api_client_connected(noise_psk=NEW_KEY.decode()) as client:
            assert (await client.device_info()).name == "zero-psk-plaintext-test"

        with pytest.raises(InvalidEncryptionKeyAPIError):
            async with api_client_connected(noise_psk=ZERO_PSK, timeout=5) as client:
                await client.device_info()