Skip to content

Commit bb44987

Browse files
arturpragaczarturpragacz
authored
Clear dynamic encryption key in ESPHome on remove (#155858)
1 parent 8d3ef2b commit bb44987

File tree

2 files changed

+194
-3
lines changed

2 files changed

+194
-3
lines changed

‎homeassistant/components/esphome/__init__.py‎

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@
22

33
from __future__ import annotations
44

5-
from aioesphomeapi import APIClient
5+
import logging
6+
7+
from aioesphomeapi import APIClient, APIConnectionError
68

79
from homeassistant.components import zeroconf
810
from homeassistant.components.bluetooth import async_remove_scanner
@@ -20,9 +22,12 @@
2022
from . import assist_satellite, dashboard, ffmpeg_proxy
2123
from .const import CONF_BLUETOOTH_MAC_ADDRESS, CONF_NOISE_PSK, DOMAIN
2224
from .domain_data import DomainData
25+
from .encryption_key_storage import async_get_encryption_key_storage
2326
from .entry_data import ESPHomeConfigEntry, RuntimeEntryData
2427
from .manager import DEVICE_CONFLICT_ISSUE_FORMAT, ESPHomeManager, cleanup_instance
2528

29+
_LOGGER = logging.getLogger(__name__)
30+
2631
CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
2732

2833
CLIENT_INFO = f"Home Assistant {ha_version}"
@@ -91,3 +96,57 @@ async def async_remove_entry(hass: HomeAssistant, entry: ESPHomeConfigEntry) ->
9196
hass, DOMAIN, DEVICE_CONFLICT_ISSUE_FORMAT.format(entry.entry_id)
9297
)
9398
await DomainData.get(hass).get_or_create_store(hass, entry).async_remove()
99+
100+
await _async_clear_dynamic_encryption_key(hass, entry)
101+
102+
103+
async def _async_clear_dynamic_encryption_key(
104+
hass: HomeAssistant, entry: ESPHomeConfigEntry
105+
) -> None:
106+
"""Clear the dynamic encryption key on the device and from storage."""
107+
if entry.unique_id is None or entry.data.get(CONF_NOISE_PSK) is None:
108+
return
109+
110+
# Only clear the key if it's stored in our storage, meaning it was
111+
# dynamically generated by us and not user-provided
112+
storage = await async_get_encryption_key_storage(hass)
113+
if await storage.async_get_key(entry.unique_id) is None:
114+
return
115+
116+
host: str = entry.data[CONF_HOST]
117+
port: int = entry.data[CONF_PORT]
118+
password: str | None = entry.data[CONF_PASSWORD]
119+
noise_psk: str | None = entry.data.get(CONF_NOISE_PSK)
120+
121+
zeroconf_instance = await zeroconf.async_get_instance(hass)
122+
123+
cli = APIClient(
124+
host,
125+
port,
126+
password,
127+
client_info=CLIENT_INFO,
128+
zeroconf_instance=zeroconf_instance,
129+
noise_psk=noise_psk,
130+
timezone=hass.config.time_zone,
131+
)
132+
133+
try:
134+
await cli.connect()
135+
# Clear the encryption key on the device by passing an empty key
136+
if not await cli.noise_encryption_set_key(b""):
137+
_LOGGER.debug(
138+
"Could not clear dynamic encryption key for ESPHome device %s: Device rejected key removal",
139+
entry.unique_id,
140+
)
141+
return
142+
except APIConnectionError as exc:
143+
_LOGGER.debug(
144+
"Could not connect to ESPHome device %s to clear dynamic encryption key: %s",
145+
entry.unique_id,
146+
exc,
147+
)
148+
return
149+
finally:
150+
await cli.disconnect()
151+
152+
await storage.async_remove_key(entry.unique_id)

‎tests/components/esphome/test_init.py‎

Lines changed: 134 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,24 @@
11
"""ESPHome set up tests."""
22

3+
from unittest.mock import AsyncMock
4+
5+
from aioesphomeapi import APIConnectionError
36
import pytest
47

58
from homeassistant.components.esphome import DOMAIN
9+
from homeassistant.components.esphome.const import CONF_NOISE_PSK
10+
from homeassistant.components.esphome.encryption_key_storage import (
11+
async_get_encryption_key_storage,
12+
)
613
from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT
714
from homeassistant.core import HomeAssistant
815

916
from tests.common import MockConfigEntry
1017

1118

1219
@pytest.mark.usefixtures("mock_client", "mock_zeroconf")
13-
async def test_delete_entry(hass: HomeAssistant) -> None:
14-
"""Test we can delete an entry without error."""
20+
async def test_remove_entry(hass: HomeAssistant) -> None:
21+
"""Test we can remove an entry without error."""
1522
entry = MockConfigEntry(
1623
domain=DOMAIN,
1724
data={CONF_HOST: "test.local", CONF_PORT: 6053, CONF_PASSWORD: ""},
@@ -22,3 +29,128 @@ async def test_delete_entry(hass: HomeAssistant) -> None:
2229
await hass.async_block_till_done()
2330
assert await hass.config_entries.async_remove(entry.entry_id)
2431
await hass.async_block_till_done()
32+
33+
34+
@pytest.mark.usefixtures("mock_zeroconf")
35+
async def test_remove_entry_clears_dynamic_encryption_key(
36+
hass: HomeAssistant,
37+
mock_client,
38+
mock_config_entry: MockConfigEntry,
39+
) -> None:
40+
"""Test that removing an entry clears the dynamic encryption key from device and storage."""
41+
# Store the encryption key to simulate it was dynamically generated
42+
storage = await async_get_encryption_key_storage(hass)
43+
await storage.async_store_key(
44+
mock_config_entry.unique_id, mock_config_entry.data[CONF_NOISE_PSK]
45+
)
46+
assert (
47+
await storage.async_get_key(mock_config_entry.unique_id)
48+
== mock_config_entry.data[CONF_NOISE_PSK]
49+
)
50+
51+
mock_client.noise_encryption_set_key = AsyncMock(return_value=True)
52+
53+
assert await hass.config_entries.async_remove(mock_config_entry.entry_id)
54+
await hass.async_block_till_done()
55+
56+
mock_client.connect.assert_called_once()
57+
mock_client.noise_encryption_set_key.assert_called_once_with(b"")
58+
mock_client.disconnect.assert_called_once()
59+
60+
assert await storage.async_get_key(mock_config_entry.unique_id) is None
61+
62+
63+
@pytest.mark.usefixtures("mock_zeroconf")
64+
async def test_remove_entry_no_noise_psk(hass: HomeAssistant, mock_client) -> None:
65+
"""Test that removing an entry without noise_psk does not attempt to clear encryption key."""
66+
mock_config_entry = MockConfigEntry(
67+
domain=DOMAIN,
68+
data={
69+
CONF_HOST: "test.local",
70+
CONF_PORT: 6053,
71+
# No CONF_NOISE_PSK
72+
},
73+
unique_id="11:22:33:44:55:aa",
74+
)
75+
mock_config_entry.add_to_hass(hass)
76+
77+
mock_client.noise_encryption_set_key = AsyncMock(return_value=True)
78+
79+
assert await hass.config_entries.async_remove(mock_config_entry.entry_id)
80+
await hass.async_block_till_done()
81+
82+
mock_client.noise_encryption_set_key.assert_not_called()
83+
84+
85+
@pytest.mark.usefixtures("mock_zeroconf")
86+
async def test_remove_entry_user_provided_key(
87+
hass: HomeAssistant,
88+
mock_client,
89+
mock_config_entry: MockConfigEntry,
90+
) -> None:
91+
"""Test that removing an entry with user-provided key does not clear it."""
92+
# Do not store the key in storage - simulates user-provided key
93+
storage = await async_get_encryption_key_storage(hass)
94+
assert await storage.async_get_key(mock_config_entry.unique_id) is None
95+
96+
mock_client.noise_encryption_set_key = AsyncMock(return_value=True)
97+
98+
assert await hass.config_entries.async_remove(mock_config_entry.entry_id)
99+
await hass.async_block_till_done()
100+
101+
mock_client.noise_encryption_set_key.assert_not_called()
102+
103+
104+
@pytest.mark.usefixtures("mock_zeroconf")
105+
async def test_remove_entry_device_rejects_key_removal(
106+
hass: HomeAssistant,
107+
mock_client,
108+
mock_config_entry: MockConfigEntry,
109+
) -> None:
110+
"""Test that when device rejects key removal, key remains in storage."""
111+
# Store the encryption key to simulate it was dynamically generated
112+
storage = await async_get_encryption_key_storage(hass)
113+
await storage.async_store_key(
114+
mock_config_entry.unique_id, mock_config_entry.data[CONF_NOISE_PSK]
115+
)
116+
117+
mock_client.noise_encryption_set_key = AsyncMock(return_value=False)
118+
119+
assert await hass.config_entries.async_remove(mock_config_entry.entry_id)
120+
await hass.async_block_till_done()
121+
122+
mock_client.connect.assert_called_once()
123+
mock_client.noise_encryption_set_key.assert_called_once_with(b"")
124+
mock_client.disconnect.assert_called_once()
125+
126+
assert (
127+
await storage.async_get_key(mock_config_entry.unique_id)
128+
== mock_config_entry.data[CONF_NOISE_PSK]
129+
)
130+
131+
132+
@pytest.mark.usefixtures("mock_zeroconf")
133+
async def test_remove_entry_connection_error(
134+
hass: HomeAssistant,
135+
mock_client,
136+
mock_config_entry: MockConfigEntry,
137+
) -> None:
138+
"""Test that connection error during key clearing does not remove key from storage."""
139+
# Store the encryption key to simulate it was dynamically generated
140+
storage = await async_get_encryption_key_storage(hass)
141+
await storage.async_store_key(
142+
mock_config_entry.unique_id, mock_config_entry.data[CONF_NOISE_PSK]
143+
)
144+
145+
mock_client.connect = AsyncMock(side_effect=APIConnectionError("Connection failed"))
146+
147+
assert await hass.config_entries.async_remove(mock_config_entry.entry_id)
148+
await hass.async_block_till_done()
149+
150+
mock_client.connect.assert_called_once()
151+
mock_client.disconnect.assert_called_once()
152+
153+
assert (
154+
await storage.async_get_key(mock_config_entry.unique_id)
155+
== mock_config_entry.data[CONF_NOISE_PSK]
156+
)

0 commit comments

Comments
 (0)