ESPHome Noise Transport Encryption support (#56216)

This commit is contained in:
Otto Winter 2021-09-20 09:02:17 +02:00 committed by GitHub
parent be19c676fa
commit a54854d129
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 341 additions and 61 deletions

View File

@ -18,6 +18,8 @@ from aioesphomeapi import (
EntityInfo,
EntityState,
HomeassistantServiceCall,
InvalidEncryptionKeyAPIError,
RequiresEncryptionAPIError,
UserService,
UserServiceArgType,
)
@ -52,6 +54,7 @@ from homeassistant.helpers.template import Template
from .entry_data import RuntimeEntryData
DOMAIN = "esphome"
CONF_NOISE_PSK = "noise_psk"
_LOGGER = logging.getLogger(__name__)
_T = TypeVar("_T")
@ -110,6 +113,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
host = entry.data[CONF_HOST]
port = entry.data[CONF_PORT]
password = entry.data[CONF_PASSWORD]
noise_psk = entry.data.get(CONF_NOISE_PSK)
device_id = None
zeroconf_instance = await zeroconf.async_get_instance(hass)
@ -121,6 +125,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
password,
client_info=f"Home Assistant {const.__version__}",
zeroconf_instance=zeroconf_instance,
noise_psk=noise_psk,
)
domain_data = DomainData.get(hass)
@ -399,6 +404,11 @@ class ReconnectLogic(RecordUpdateListener):
try:
await self._cli.connect(on_stop=self._on_disconnect, login=True)
except APIConnectionError as error:
if isinstance(
error, (RequiresEncryptionAPIError, InvalidEncryptionKeyAPIError)
):
self._entry.async_start_reauth(self._hass)
level = logging.WARNING if tries == 0 else logging.DEBUG
_LOGGER.log(
level,

View File

@ -4,7 +4,15 @@ from __future__ import annotations
from collections import OrderedDict
from typing import Any
from aioesphomeapi import APIClient, APIConnectionError, DeviceInfo
from aioesphomeapi import (
APIClient,
APIConnectionError,
DeviceInfo,
InvalidAuthAPIError,
InvalidEncryptionKeyAPIError,
RequiresEncryptionAPIError,
ResolveAPIError,
)
import voluptuous as vol
from homeassistant.components import zeroconf
@ -14,7 +22,9 @@ from homeassistant.core import callback
from homeassistant.data_entry_flow import FlowResult
from homeassistant.helpers.typing import DiscoveryInfoType
from . import DOMAIN, DomainData
from . import CONF_NOISE_PSK, DOMAIN, DomainData
ERROR_REQUIRES_ENCRYPTION_KEY = "requires_encryption_key"
class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
@ -27,12 +37,14 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
self._host: str | None = None
self._port: int | None = None
self._password: str | None = None
self._noise_psk: str | None = None
self._device_info: DeviceInfo | None = None
async def _async_step_user_base(
self, user_input: dict[str, Any] | None = None, error: str | None = None
) -> FlowResult:
if user_input is not None:
return await self._async_authenticate_or_add(user_input)
return await self._async_try_fetch_device_info(user_input)
fields: dict[Any, type] = OrderedDict()
fields[vol.Required(CONF_HOST, default=self._host or vol.UNDEFINED)] = str
@ -52,6 +64,36 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
"""Handle a flow initialized by the user."""
return await self._async_step_user_base(user_input=user_input)
async def async_step_reauth(self, data: dict[str, Any] | None = None) -> FlowResult:
"""Handle a flow initialized by a reauth event."""
entry = self.hass.config_entries.async_get_entry(self.context["entry_id"])
assert entry is not None
self._host = entry.data[CONF_HOST]
self._port = entry.data[CONF_PORT]
self._password = entry.data[CONF_PASSWORD]
self._noise_psk = entry.data.get(CONF_NOISE_PSK)
return await self.async_step_reauth_confirm()
async def async_step_reauth_confirm(
self, user_input: dict[str, Any] | None = None
) -> FlowResult:
"""Handle reauthorization flow."""
errors = {}
if user_input is not None:
self._noise_psk = user_input[CONF_NOISE_PSK]
error = await self.fetch_device_info()
if error is None:
return await self._async_authenticate_or_add()
errors["base"] = error
return self.async_show_form(
step_id="reauth_confirm",
data_schema=vol.Schema({vol.Required(CONF_NOISE_PSK): str}),
errors=errors,
description_placeholders={"name": self._name},
)
@property
def _name(self) -> str | None:
return self.context.get(CONF_NAME)
@ -67,18 +109,21 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
self._host = user_input[CONF_HOST]
self._port = user_input[CONF_PORT]
async def _async_authenticate_or_add(
async def _async_try_fetch_device_info(
self, user_input: dict[str, Any] | None
) -> FlowResult:
self._set_user_input(user_input)
error, device_info = await self.fetch_device_info()
error = await self.fetch_device_info()
if error == ERROR_REQUIRES_ENCRYPTION_KEY:
return await self.async_step_encryption_key()
if error is not None:
return await self._async_step_user_base(error=error)
assert device_info is not None
self._name = device_info.name
return await self._async_authenticate_or_add()
async def _async_authenticate_or_add(self) -> FlowResult:
# Only show authentication step if device uses password
if device_info.uses_password:
assert self._device_info is not None
if self._device_info.uses_password:
return await self.async_step_authenticate()
return self._async_get_entry()
@ -88,7 +133,7 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
) -> FlowResult:
"""Handle user-confirmation of discovered node."""
if user_input is not None:
return await self._async_authenticate_or_add(None)
return await self._async_try_fetch_device_info(None)
return self.async_show_form(
step_id="discovery_confirm", description_placeholders={"name": self._name}
)
@ -144,15 +189,47 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
@callback
def _async_get_entry(self) -> FlowResult:
assert self._name is not None
return self.async_create_entry(
title=self._name,
data={
config_data = {
CONF_HOST: self._host,
CONF_PORT: self._port,
# The API uses protobuf, so empty string denotes absence
CONF_PASSWORD: self._password or "",
},
CONF_NOISE_PSK: self._noise_psk or "",
}
if "entry_id" in self.context:
entry = self.hass.config_entries.async_get_entry(self.context["entry_id"])
assert entry is not None
self.hass.config_entries.async_update_entry(entry, data=config_data)
# Reload the config entry to notify of updated config
self.hass.async_create_task(
self.hass.config_entries.async_reload(entry.entry_id)
)
return self.async_abort(reason="reauth_successful")
assert self._name is not None
return self.async_create_entry(
title=self._name,
data=config_data,
)
async def async_step_encryption_key(
self, user_input: dict[str, Any] | None = None
) -> FlowResult:
"""Handle getting psk for transport encryption."""
errors = {}
if user_input is not None:
self._noise_psk = user_input[CONF_NOISE_PSK]
error = await self.fetch_device_info()
if error is None:
return await self._async_authenticate_or_add()
errors["base"] = error
return self.async_show_form(
step_id="encryption_key",
data_schema=vol.Schema({vol.Required(CONF_NOISE_PSK): str}),
errors=errors,
description_placeholders={"name": self._name},
)
async def async_step_authenticate(
@ -177,7 +254,7 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
errors=errors,
)
async def fetch_device_info(self) -> tuple[str | None, DeviceInfo | None]:
async def fetch_device_info(self) -> str | None:
"""Fetch device info from API and return any errors."""
zeroconf_instance = await zeroconf.async_get_instance(self.hass)
assert self._host is not None
@ -188,19 +265,26 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
self._port,
"",
zeroconf_instance=zeroconf_instance,
noise_psk=self._noise_psk,
)
try:
await cli.connect()
device_info = await cli.device_info()
except APIConnectionError as err:
if "resolving" in str(err):
return "resolve_error", None
return "connection_error", None
self._device_info = await cli.device_info()
except RequiresEncryptionAPIError:
return ERROR_REQUIRES_ENCRYPTION_KEY
except InvalidEncryptionKeyAPIError:
return "invalid_psk"
except ResolveAPIError:
return "resolve_error"
except APIConnectionError:
return "connection_error"
finally:
await cli.disconnect(force=True)
return None, device_info
self._name = self._device_info.name
return None
async def try_login(self) -> str | None:
"""Try logging in to device and return any errors."""
@ -213,12 +297,16 @@ class EsphomeFlowHandler(ConfigFlow, domain=DOMAIN):
self._port,
self._password,
zeroconf_instance=zeroconf_instance,
noise_psk=self._noise_psk,
)
try:
await cli.connect(login=True)
except APIConnectionError:
await cli.disconnect(force=True)
except InvalidAuthAPIError:
return "invalid_auth"
except APIConnectionError:
return "connection_error"
finally:
await cli.disconnect(force=True)
return None

View File

@ -3,7 +3,7 @@
"name": "ESPHome",
"config_flow": true,
"documentation": "https://www.home-assistant.io/integrations/esphome",
"requirements": ["aioesphomeapi==8.0.0"],
"requirements": ["aioesphomeapi==9.1.0"],
"zeroconf": ["_esphomelib._tcp.local."],
"codeowners": ["@OttoWinter", "@jesserockz"],
"after_dependencies": ["zeroconf", "tag"],

View File

@ -2,12 +2,14 @@
"config": {
"abort": {
"already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
"already_in_progress": "[%key:common::config_flow::abort::already_in_progress%]"
"already_in_progress": "[%key:common::config_flow::abort::already_in_progress%]",
"reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
},
"error": {
"resolve_error": "Can't resolve address of the ESP. If this error persists, please set a static IP address: https://esphomelib.com/esphomeyaml/components/wifi.html#manual-ips",
"connection_error": "Can't connect to ESP. Please make sure your YAML file contains an 'api:' line.",
"invalid_auth": "[%key:common::config_flow::error::invalid_auth%]"
"invalid_auth": "[%key:common::config_flow::error::invalid_auth%]",
"invalid_psk": "The transport encryption key is invalid. Please ensure it matches what you have in your configuration"
},
"step": {
"user": {
@ -23,6 +25,18 @@
},
"description": "Please enter the password you set in your configuration for {name}."
},
"encryption_key": {
"data": {
"noise_psk": "Encryption key"
},
"description": "Please enter the encryption key you set in your configuration for {name}."
},
"reauth_confirm": {
"data": {
"noise_psk": "Encryption key"
},
"description": "The ESPHome device {name} enabled transport encryption or changed the encryption key. Please enter the updated key."
},
"discovery_confirm": {
"description": "Do you want to add the ESPHome node `{name}` to Home Assistant?",
"title": "Discovered ESPHome node"

View File

@ -164,7 +164,7 @@ aioeagle==1.1.0
aioemonitor==1.0.5
# homeassistant.components.esphome
aioesphomeapi==8.0.0
aioesphomeapi==9.1.0
# homeassistant.components.flo
aioflo==0.4.1

View File

@ -109,7 +109,7 @@ aioeagle==1.1.0
aioemonitor==1.0.5
# homeassistant.components.esphome
aioesphomeapi==8.0.0
aioesphomeapi==9.1.0
# homeassistant.components.flo
aioflo==0.4.1

View File

@ -2,10 +2,17 @@
from collections import namedtuple
from unittest.mock import AsyncMock, MagicMock, patch
from aioesphomeapi import (
APIConnectionError,
InvalidAuthAPIError,
InvalidEncryptionKeyAPIError,
RequiresEncryptionAPIError,
ResolveAPIError,
)
import pytest
from homeassistant import config_entries
from homeassistant.components.esphome import DOMAIN, DomainData
from homeassistant.components.esphome import CONF_NOISE_PSK, DOMAIN, DomainData
from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT
from homeassistant.data_entry_flow import (
RESULT_TYPE_ABORT,
@ -16,6 +23,8 @@ from homeassistant.data_entry_flow import (
from tests.common import MockConfigEntry
MockDeviceInfo = namedtuple("DeviceInfo", ["uses_password", "name"])
VALID_NOISE_PSK = "bOFFzzvfpg5DB94DuBGLXD/hMnhpDKgP9UQyBulwWVU="
INVALID_NOISE_PSK = "lSYBYEjQI1bVL8s2Vask4YytGMj1f1epNtmoim2yuTM="
@pytest.fixture
@ -23,12 +32,15 @@ def mock_client():
"""Mock APIClient."""
with patch("homeassistant.components.esphome.config_flow.APIClient") as mock_client:
def mock_constructor(loop, host, port, password, zeroconf_instance=None):
def mock_constructor(
loop, host, port, password, zeroconf_instance=None, noise_psk=None
):
"""Fake the client constructor."""
mock_client.host = host
mock_client.port = port
mock_client.password = password
mock_client.zeroconf_instance = zeroconf_instance
mock_client.noise_psk = noise_psk
return mock_client
mock_client.side_effect = mock_constructor
@ -38,16 +50,6 @@ def mock_client():
yield mock_client
@pytest.fixture(autouse=True)
def mock_api_connection_error():
"""Mock out the try login method."""
with patch(
"homeassistant.components.esphome.config_flow.APIConnectionError",
new_callable=lambda: OSError,
) as mock_error:
yield mock_error
@pytest.fixture(autouse=True)
def mock_setup_entry():
"""Mock setting up a config entry."""
@ -75,7 +77,12 @@ async def test_user_connection_works(hass, mock_client, mock_zeroconf):
)
assert result["type"] == RESULT_TYPE_CREATE_ENTRY
assert result["data"] == {CONF_HOST: "127.0.0.1", CONF_PORT: 80, CONF_PASSWORD: ""}
assert result["data"] == {
CONF_HOST: "127.0.0.1",
CONF_PORT: 80,
CONF_PASSWORD: "",
CONF_NOISE_PSK: "",
}
assert result["title"] == "test"
assert len(mock_client.connect.mock_calls) == 1
@ -84,23 +91,15 @@ async def test_user_connection_works(hass, mock_client, mock_zeroconf):
assert mock_client.host == "127.0.0.1"
assert mock_client.port == 80
assert mock_client.password == ""
assert mock_client.noise_psk is None
async def test_user_resolve_error(
hass, mock_api_connection_error, mock_client, mock_zeroconf
):
async def test_user_resolve_error(hass, mock_client, mock_zeroconf):
"""Test user step with IP resolve error."""
class MockResolveError(mock_api_connection_error):
"""Create an exception with a specific error message."""
def __init__(self):
"""Initialize."""
super().__init__("Error resolving IP address")
with patch(
"homeassistant.components.esphome.config_flow.APIConnectionError",
new_callable=lambda: MockResolveError,
new_callable=lambda: ResolveAPIError,
) as exc:
mock_client.device_info.side_effect = exc
result = await hass.config_entries.flow.async_init(
@ -118,11 +117,9 @@ async def test_user_resolve_error(
assert len(mock_client.disconnect.mock_calls) == 1
async def test_user_connection_error(
hass, mock_api_connection_error, mock_client, mock_zeroconf
):
async def test_user_connection_error(hass, mock_client, mock_zeroconf):
"""Test user step with connection error."""
mock_client.device_info.side_effect = mock_api_connection_error
mock_client.device_info.side_effect = APIConnectionError
result = await hass.config_entries.flow.async_init(
"esphome",
@ -161,13 +158,12 @@ async def test_user_with_password(hass, mock_client, mock_zeroconf):
CONF_HOST: "127.0.0.1",
CONF_PORT: 6053,
CONF_PASSWORD: "password1",
CONF_NOISE_PSK: "",
}
assert mock_client.password == "password1"
async def test_user_invalid_password(
hass, mock_api_connection_error, mock_client, mock_zeroconf
):
async def test_user_invalid_password(hass, mock_client, mock_zeroconf):
"""Test user step with invalid password."""
mock_client.device_info = AsyncMock(return_value=MockDeviceInfo(True, "test"))
@ -180,7 +176,7 @@ async def test_user_invalid_password(
assert result["type"] == RESULT_TYPE_FORM
assert result["step_id"] == "authenticate"
mock_client.connect.side_effect = mock_api_connection_error
mock_client.connect.side_effect = InvalidAuthAPIError
result = await hass.config_entries.flow.async_configure(
result["flow_id"], user_input={CONF_PASSWORD: "invalid"}
@ -191,6 +187,30 @@ async def test_user_invalid_password(
assert result["errors"] == {"base": "invalid_auth"}
async def test_login_connection_error(hass, mock_client, mock_zeroconf):
"""Test user step with connection error on login attempt."""
mock_client.device_info = AsyncMock(return_value=MockDeviceInfo(True, "test"))
result = await hass.config_entries.flow.async_init(
"esphome",
context={"source": config_entries.SOURCE_USER},
data={CONF_HOST: "127.0.0.1", CONF_PORT: 6053},
)
assert result["type"] == RESULT_TYPE_FORM
assert result["step_id"] == "authenticate"
mock_client.connect.side_effect = APIConnectionError
result = await hass.config_entries.flow.async_configure(
result["flow_id"], user_input={CONF_PASSWORD: "valid"}
)
assert result["type"] == RESULT_TYPE_FORM
assert result["step_id"] == "authenticate"
assert result["errors"] == {"base": "connection_error"}
async def test_discovery_initiation(hass, mock_client, mock_zeroconf):
"""Test discovery importing works."""
mock_client.device_info = AsyncMock(return_value=MockDeviceInfo(False, "test8266"))
@ -345,3 +365,151 @@ async def test_discovery_updates_unique_id(hass, mock_client):
assert result["reason"] == "already_configured"
assert entry.unique_id == "test8266"
async def test_user_requires_psk(hass, mock_client, mock_zeroconf):
"""Test user step with requiring encryption key."""
mock_client.device_info.side_effect = RequiresEncryptionAPIError
result = await hass.config_entries.flow.async_init(
"esphome",
context={"source": config_entries.SOURCE_USER},
data={CONF_HOST: "127.0.0.1", CONF_PORT: 6053},
)
assert result["type"] == RESULT_TYPE_FORM
assert result["step_id"] == "encryption_key"
assert result["errors"] == {}
assert len(mock_client.connect.mock_calls) == 1
assert len(mock_client.device_info.mock_calls) == 1
assert len(mock_client.disconnect.mock_calls) == 1
async def test_encryption_key_valid_psk(hass, mock_client, mock_zeroconf):
"""Test encryption key step with valid key."""
mock_client.device_info.side_effect = RequiresEncryptionAPIError
result = await hass.config_entries.flow.async_init(
"esphome",
context={"source": config_entries.SOURCE_USER},
data={CONF_HOST: "127.0.0.1", CONF_PORT: 6053},
)
assert result["type"] == RESULT_TYPE_FORM
assert result["step_id"] == "encryption_key"
mock_client.device_info = AsyncMock(return_value=MockDeviceInfo(False, "test"))
result = await hass.config_entries.flow.async_configure(
result["flow_id"], user_input={CONF_NOISE_PSK: VALID_NOISE_PSK}
)
assert result["type"] == RESULT_TYPE_CREATE_ENTRY
assert result["data"] == {
CONF_HOST: "127.0.0.1",
CONF_PORT: 6053,
CONF_PASSWORD: "",
CONF_NOISE_PSK: VALID_NOISE_PSK,
}
assert mock_client.noise_psk == VALID_NOISE_PSK
async def test_encryption_key_invalid_psk(hass, mock_client, mock_zeroconf):
"""Test encryption key step with invalid key."""
mock_client.device_info.side_effect = RequiresEncryptionAPIError
result = await hass.config_entries.flow.async_init(
"esphome",
context={"source": config_entries.SOURCE_USER},
data={CONF_HOST: "127.0.0.1", CONF_PORT: 6053},
)
assert result["type"] == RESULT_TYPE_FORM
assert result["step_id"] == "encryption_key"
mock_client.device_info.side_effect = InvalidEncryptionKeyAPIError
result = await hass.config_entries.flow.async_configure(
result["flow_id"], user_input={CONF_NOISE_PSK: INVALID_NOISE_PSK}
)
assert result["type"] == RESULT_TYPE_FORM
assert result["step_id"] == "encryption_key"
assert result["errors"] == {"base": "invalid_psk"}
assert mock_client.noise_psk == INVALID_NOISE_PSK
async def test_reauth_initiation(hass, mock_client, mock_zeroconf):
"""Test reauth initiation shows form."""
entry = MockConfigEntry(
domain=DOMAIN,
data={CONF_HOST: "127.0.0.1", CONF_PORT: 6053, CONF_PASSWORD: ""},
)
entry.add_to_hass(hass)
result = await hass.config_entries.flow.async_init(
"esphome",
context={
"source": config_entries.SOURCE_REAUTH,
"entry_id": entry.entry_id,
"unique_id": entry.unique_id,
},
)
assert result["type"] == RESULT_TYPE_FORM
assert result["step_id"] == "reauth_confirm"
async def test_reauth_confirm_valid(hass, mock_client, mock_zeroconf):
"""Test reauth initiation with valid PSK."""
entry = MockConfigEntry(
domain=DOMAIN,
data={CONF_HOST: "127.0.0.1", CONF_PORT: 6053, CONF_PASSWORD: ""},
)
entry.add_to_hass(hass)
result = await hass.config_entries.flow.async_init(
"esphome",
context={
"source": config_entries.SOURCE_REAUTH,
"entry_id": entry.entry_id,
"unique_id": entry.unique_id,
},
)
mock_client.device_info = AsyncMock(return_value=MockDeviceInfo(False, "test"))
result = await hass.config_entries.flow.async_configure(
result["flow_id"], user_input={CONF_NOISE_PSK: VALID_NOISE_PSK}
)
assert result["type"] == RESULT_TYPE_ABORT
assert result["reason"] == "reauth_successful"
assert entry.data[CONF_NOISE_PSK] == VALID_NOISE_PSK
async def test_reauth_confirm_invalid(hass, mock_client, mock_zeroconf):
"""Test reauth initiation with invalid PSK."""
entry = MockConfigEntry(
domain=DOMAIN,
data={CONF_HOST: "127.0.0.1", CONF_PORT: 6053, CONF_PASSWORD: ""},
)
entry.add_to_hass(hass)
result = await hass.config_entries.flow.async_init(
"esphome",
context={
"source": config_entries.SOURCE_REAUTH,
"entry_id": entry.entry_id,
"unique_id": entry.unique_id,
},
)
mock_client.device_info.side_effect = InvalidEncryptionKeyAPIError
result = await hass.config_entries.flow.async_configure(
result["flow_id"], user_input={CONF_NOISE_PSK: INVALID_NOISE_PSK}
)
assert result["type"] == RESULT_TYPE_FORM
assert result["step_id"] == "reauth_confirm"
assert result["errors"]
assert result["errors"]["base"] == "invalid_psk"