Add Multi factor authentication support for Sense (#66498)

Co-authored-by: J. Nick Koston <nick@koston.org>
This commit is contained in:
Keilin Bickar 2022-02-21 17:05:12 -05:00 committed by GitHub
parent ba2bc975f4
commit e6af7847fc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 359 additions and 57 deletions

View File

@ -2,7 +2,7 @@
"domain": "emulated_kasa", "domain": "emulated_kasa",
"name": "Emulated Kasa", "name": "Emulated Kasa",
"documentation": "https://www.home-assistant.io/integrations/emulated_kasa", "documentation": "https://www.home-assistant.io/integrations/emulated_kasa",
"requirements": ["sense_energy==0.9.6"], "requirements": ["sense_energy==0.10.2"],
"codeowners": ["@kbickar"], "codeowners": ["@kbickar"],
"quality_scale": "internal", "quality_scale": "internal",
"iot_class": "local_push", "iot_class": "local_push",

View File

@ -3,18 +3,21 @@ import asyncio
from datetime import timedelta from datetime import timedelta
import logging import logging
from sense_energy import ASyncSenseable, SenseAuthenticationException from sense_energy import (
ASyncSenseable,
SenseAuthenticationException,
SenseMFARequiredException,
)
from homeassistant.config_entries import ConfigEntry from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ( from homeassistant.const import (
CONF_EMAIL, CONF_EMAIL,
CONF_PASSWORD,
CONF_TIMEOUT, CONF_TIMEOUT,
EVENT_HOMEASSISTANT_STOP, EVENT_HOMEASSISTANT_STOP,
Platform, Platform,
) )
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.aiohttp_client import async_get_clientsession
from homeassistant.helpers.dispatcher import async_dispatcher_send from homeassistant.helpers.dispatcher import async_dispatcher_send
from homeassistant.helpers.event import async_track_time_interval from homeassistant.helpers.event import async_track_time_interval
@ -58,9 +61,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
entry_data = entry.data entry_data = entry.data
email = entry_data[CONF_EMAIL] email = entry_data[CONF_EMAIL]
password = entry_data[CONF_PASSWORD]
timeout = entry_data[CONF_TIMEOUT] timeout = entry_data[CONF_TIMEOUT]
access_token = entry_data.get("access_token", "")
user_id = entry_data.get("user_id", "")
monitor_id = entry_data.get("monitor_id", "")
client_session = async_get_clientsession(hass) client_session = async_get_clientsession(hass)
gateway = ASyncSenseable( gateway = ASyncSenseable(
@ -69,16 +75,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
gateway.rate_limit = ACTIVE_UPDATE_RATE gateway.rate_limit = ACTIVE_UPDATE_RATE
try: try:
await gateway.authenticate(email, password) gateway.load_auth(access_token, user_id, monitor_id)
except SenseAuthenticationException: await gateway.get_monitor_data()
_LOGGER.error("Could not authenticate with sense server") except (SenseAuthenticationException, SenseMFARequiredException) as err:
return False _LOGGER.warning("Sense authentication expired")
except SENSE_TIMEOUT_EXCEPTIONS as err: raise ConfigEntryAuthFailed(err) from err
raise ConfigEntryNotReady(
str(err) or "Timed out during authentication"
) from err
except SENSE_EXCEPTIONS as err:
raise ConfigEntryNotReady(str(err) or "Error during authentication") from err
sense_devices_data = SenseDevicesData() sense_devices_data = SenseDevicesData()
try: try:
@ -91,11 +92,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
except SENSE_EXCEPTIONS as err: except SENSE_EXCEPTIONS as err:
raise ConfigEntryNotReady(str(err) or "Error during realtime update") from err raise ConfigEntryNotReady(str(err) or "Error during realtime update") from err
async def _async_update_trend():
"""Update the trend data."""
try:
await gateway.update_trend_data()
except (SenseAuthenticationException, SenseMFARequiredException) as err:
_LOGGER.warning("Sense authentication expired")
raise ConfigEntryAuthFailed(err) from err
trends_coordinator: DataUpdateCoordinator[None] = DataUpdateCoordinator( trends_coordinator: DataUpdateCoordinator[None] = DataUpdateCoordinator(
hass, hass,
_LOGGER, _LOGGER,
name=f"Sense Trends {email}", name=f"Sense Trends {email}",
update_method=gateway.update_trend_data, update_method=_async_update_trend,
update_interval=timedelta(seconds=300), update_interval=timedelta(seconds=300),
) )
# Start out as unavailable so we do not report 0 data # Start out as unavailable so we do not report 0 data

View File

@ -1,11 +1,15 @@
"""Config flow for Sense integration.""" """Config flow for Sense integration."""
import logging import logging
from sense_energy import ASyncSenseable, SenseAuthenticationException from sense_energy import (
ASyncSenseable,
SenseAuthenticationException,
SenseMFARequiredException,
)
import voluptuous as vol import voluptuous as vol
from homeassistant import config_entries, core from homeassistant import config_entries
from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, CONF_TIMEOUT from homeassistant.const import CONF_CODE, CONF_EMAIL, CONF_PASSWORD, CONF_TIMEOUT
from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.aiohttp_client import async_get_clientsession
from .const import ACTIVE_UPDATE_RATE, DEFAULT_TIMEOUT, DOMAIN, SENSE_TIMEOUT_EXCEPTIONS from .const import ACTIVE_UPDATE_RATE, DEFAULT_TIMEOUT, DOMAIN, SENSE_TIMEOUT_EXCEPTIONS
@ -21,37 +25,74 @@ DATA_SCHEMA = vol.Schema(
) )
async def validate_input(hass: core.HomeAssistant, data):
"""Validate the user input allows us to connect.
Data has the keys from DATA_SCHEMA with values provided by the user.
"""
timeout = data[CONF_TIMEOUT]
client_session = async_get_clientsession(hass)
gateway = ASyncSenseable(
api_timeout=timeout, wss_timeout=timeout, client_session=client_session
)
gateway.rate_limit = ACTIVE_UPDATE_RATE
await gateway.authenticate(data[CONF_EMAIL], data[CONF_PASSWORD])
# Return info that you want to store in the config entry.
return {"title": data[CONF_EMAIL]}
class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
"""Handle a config flow for Sense.""" """Handle a config flow for Sense."""
VERSION = 1 VERSION = 1
async def async_step_user(self, user_input=None): def __init__(self):
"""Handle the initial step.""" """Init Config ."""
self._gateway = None
self._auth_data = {}
super().__init__()
async def validate_input(self, data):
"""Validate the user input allows us to connect.
Data has the keys from DATA_SCHEMA with values provided by the user.
"""
self._auth_data.update(dict(data))
timeout = self._auth_data[CONF_TIMEOUT]
client_session = async_get_clientsession(self.hass)
self._gateway = ASyncSenseable(
api_timeout=timeout, wss_timeout=timeout, client_session=client_session
)
self._gateway.rate_limit = ACTIVE_UPDATE_RATE
await self._gateway.authenticate(
self._auth_data[CONF_EMAIL], self._auth_data[CONF_PASSWORD]
)
async def create_entry_from_data(self):
"""Create the entry from the config data."""
self._auth_data["access_token"] = self._gateway.sense_access_token
self._auth_data["user_id"] = self._gateway.sense_user_id
self._auth_data["monitor_id"] = self._gateway.sense_monitor_id
existing_entry = await self.async_set_unique_id(self._auth_data[CONF_EMAIL])
if not existing_entry:
return self.async_create_entry(
title=self._auth_data[CONF_EMAIL], data=self._auth_data
)
self.hass.config_entries.async_update_entry(
existing_entry, data=self._auth_data
)
await self.hass.config_entries.async_reload(existing_entry.entry_id)
return self.async_abort(reason="reauth_successful")
async def validate_input_and_create_entry(self, user_input, errors):
"""Validate the input and create the entry from the data."""
try:
await self.validate_input(user_input)
except SenseMFARequiredException:
return await self.async_step_validation()
except SENSE_TIMEOUT_EXCEPTIONS:
errors["base"] = "cannot_connect"
except SenseAuthenticationException:
errors["base"] = "invalid_auth"
except Exception: # pylint: disable=broad-except
_LOGGER.exception("Unexpected exception")
errors["base"] = "unknown"
else:
return await self.create_entry_from_data()
return None
async def async_step_validation(self, user_input=None):
"""Handle validation (2fa) step."""
errors = {} errors = {}
if user_input is not None: if user_input:
try: try:
info = await validate_input(self.hass, user_input) await self._gateway.validate_mfa(user_input[CONF_CODE])
await self.async_set_unique_id(user_input[CONF_EMAIL])
return self.async_create_entry(title=info["title"], data=user_input)
except SENSE_TIMEOUT_EXCEPTIONS: except SENSE_TIMEOUT_EXCEPTIONS:
errors["base"] = "cannot_connect" errors["base"] = "cannot_connect"
except SenseAuthenticationException: except SenseAuthenticationException:
@ -59,7 +100,43 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
except Exception: # pylint: disable=broad-except except Exception: # pylint: disable=broad-except
_LOGGER.exception("Unexpected exception") _LOGGER.exception("Unexpected exception")
errors["base"] = "unknown" errors["base"] = "unknown"
else:
return await self.create_entry_from_data()
return self.async_show_form(
step_id="validation",
data_schema=vol.Schema({vol.Required(CONF_CODE): vol.All(str, vol.Strip)}),
errors=errors,
)
async def async_step_user(self, user_input=None):
"""Handle the initial step."""
errors = {}
if user_input is not None:
if result := await self.validate_input_and_create_entry(user_input, errors):
return result
return self.async_show_form( return self.async_show_form(
step_id="user", data_schema=DATA_SCHEMA, errors=errors step_id="user", data_schema=DATA_SCHEMA, errors=errors
) )
async def async_step_reauth(self, data):
"""Handle configuration by re-auth."""
self._auth_data = dict(data)
return await self.async_step_reauth_validate(data)
async def async_step_reauth_validate(self, user_input=None):
"""Handle reauth and validation."""
errors = {}
if user_input is not None:
if result := await self.validate_input_and_create_entry(user_input, errors):
return result
return self.async_show_form(
step_id="reauth_validate",
data_schema=vol.Schema({vol.Required(CONF_PASSWORD): str}),
errors=errors,
description_placeholders={
CONF_EMAIL: self._auth_data[CONF_EMAIL],
},
)

View File

@ -2,7 +2,7 @@
"domain": "sense", "domain": "sense",
"name": "Sense", "name": "Sense",
"documentation": "https://www.home-assistant.io/integrations/sense", "documentation": "https://www.home-assistant.io/integrations/sense",
"requirements": ["sense_energy==0.9.6"], "requirements": ["sense_energy==0.10.2"],
"codeowners": ["@kbickar"], "codeowners": ["@kbickar"],
"config_flow": true, "config_flow": true,
"dhcp": [ "dhcp": [

View File

@ -8,6 +8,19 @@
"password": "[%key:common::config_flow::data::password%]", "password": "[%key:common::config_flow::data::password%]",
"timeout": "Timeout" "timeout": "Timeout"
} }
},
"validation": {
"title": "Sense Multi-factor authentication",
"data": {
"code": "Verification code"
}
},
"reauth_validate": {
"title": "[%key:common::config_flow::title::reauth%]",
"description": "The Sense integration needs to re-authenticate your account {email}.",
"data": {
"password": "[%key:common::config_flow::data::password%]"
}
} }
}, },
"error": { "error": {
@ -16,7 +29,8 @@
"unknown": "[%key:common::config_flow::error::unknown%]" "unknown": "[%key:common::config_flow::error::unknown%]"
}, },
"abort": { "abort": {
"already_configured": "[%key:common::config_flow::abort::already_configured_device%]" "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
"reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
} }
} }
} }

View File

@ -1,7 +1,8 @@
{ {
"config": { "config": {
"abort": { "abort": {
"already_configured": "Device is already configured" "already_configured": "Device is already configured",
"reauth_successful": "Re-authentication was successful"
}, },
"error": { "error": {
"cannot_connect": "Failed to connect", "cannot_connect": "Failed to connect",
@ -9,6 +10,13 @@
"unknown": "Unexpected error" "unknown": "Unexpected error"
}, },
"step": { "step": {
"reauth_validate": {
"data": {
"password": "Password"
},
"description": "The Sense integration needs to re-authenticate your account {email}.",
"title": "Reauthenticate Integration"
},
"user": { "user": {
"data": { "data": {
"email": "Email", "email": "Email",
@ -16,6 +24,12 @@
"timeout": "Timeout" "timeout": "Timeout"
}, },
"title": "Connect to your Sense Energy Monitor" "title": "Connect to your Sense Energy Monitor"
},
"validation": {
"data": {
"code": "Verification code"
},
"title": "Sense Multi-factor authentication"
} }
} }
} }

View File

@ -2178,7 +2178,7 @@ sense-hat==2.2.0
# homeassistant.components.emulated_kasa # homeassistant.components.emulated_kasa
# homeassistant.components.sense # homeassistant.components.sense
sense_energy==0.9.6 sense_energy==0.10.2
# homeassistant.components.sentry # homeassistant.components.sentry
sentry-sdk==1.5.5 sentry-sdk==1.5.5

View File

@ -1346,7 +1346,7 @@ screenlogicpy==0.5.4
# homeassistant.components.emulated_kasa # homeassistant.components.emulated_kasa
# homeassistant.components.sense # homeassistant.components.sense
sense_energy==0.9.6 sense_energy==0.10.2
# homeassistant.components.sentry # homeassistant.components.sentry
sentry-sdk==1.5.5 sentry-sdk==1.5.5

View File

@ -1,13 +1,44 @@
"""Test the Sense config flow.""" """Test the Sense config flow."""
from unittest.mock import patch from unittest.mock import AsyncMock, patch
from sense_energy import SenseAPITimeoutException, SenseAuthenticationException import pytest
from sense_energy import (
SenseAPITimeoutException,
SenseAuthenticationException,
SenseMFARequiredException,
)
from homeassistant import config_entries from homeassistant import config_entries
from homeassistant.components.sense.const import DOMAIN from homeassistant.components.sense.const import DOMAIN
from homeassistant.const import CONF_CODE
from tests.common import MockConfigEntry
MOCK_CONFIG = {
"timeout": 6,
"email": "test-email",
"password": "test-password",
"access_token": "ABC",
"user_id": "123",
"monitor_id": "456",
}
async def test_form(hass): @pytest.fixture(name="mock_sense")
def mock_sense():
"""Mock Sense object for authenticatation."""
with patch(
"homeassistant.components.sense.config_flow.ASyncSenseable"
) as mock_sense:
mock_sense.return_value.authenticate = AsyncMock(return_value=True)
mock_sense.return_value.validate_mfa = AsyncMock(return_value=True)
mock_sense.return_value.sense_access_token = "ABC"
mock_sense.return_value.sense_user_id = "123"
mock_sense.return_value.sense_monitor_id = "456"
yield mock_sense
async def test_form(hass, mock_sense):
"""Test we get the form.""" """Test we get the form."""
result = await hass.config_entries.flow.async_init( result = await hass.config_entries.flow.async_init(
@ -16,7 +47,7 @@ async def test_form(hass):
assert result["type"] == "form" assert result["type"] == "form"
assert result["errors"] == {} assert result["errors"] == {}
with patch("sense_energy.ASyncSenseable.authenticate", return_value=True,), patch( with patch(
"homeassistant.components.sense.async_setup_entry", "homeassistant.components.sense.async_setup_entry",
return_value=True, return_value=True,
) as mock_setup_entry: ) as mock_setup_entry:
@ -28,11 +59,7 @@ async def test_form(hass):
assert result2["type"] == "create_entry" assert result2["type"] == "create_entry"
assert result2["title"] == "test-email" assert result2["title"] == "test-email"
assert result2["data"] == { assert result2["data"] == MOCK_CONFIG
"timeout": 6,
"email": "test-email",
"password": "test-password",
}
assert len(mock_setup_entry.mock_calls) == 1 assert len(mock_setup_entry.mock_calls) == 1
@ -55,6 +82,113 @@ async def test_form_invalid_auth(hass):
assert result2["errors"] == {"base": "invalid_auth"} assert result2["errors"] == {"base": "invalid_auth"}
async def test_form_mfa_required(hass, mock_sense):
"""Test we handle invalid auth."""
result = await hass.config_entries.flow.async_init(
DOMAIN, context={"source": config_entries.SOURCE_USER}
)
mock_sense.return_value.authenticate.side_effect = SenseMFARequiredException
result2 = await hass.config_entries.flow.async_configure(
result["flow_id"],
{"timeout": "6", "email": "test-email", "password": "test-password"},
)
assert result2["type"] == "form"
assert result2["step_id"] == "validation"
mock_sense.return_value.validate_mfa.side_effect = None
result3 = await hass.config_entries.flow.async_configure(
result["flow_id"],
{CONF_CODE: "012345"},
)
assert result3["type"] == "create_entry"
assert result3["title"] == "test-email"
assert result3["data"] == MOCK_CONFIG
async def test_form_mfa_required_wrong(hass, mock_sense):
"""Test we handle invalid auth."""
result = await hass.config_entries.flow.async_init(
DOMAIN, context={"source": config_entries.SOURCE_USER}
)
mock_sense.return_value.authenticate.side_effect = SenseMFARequiredException
result2 = await hass.config_entries.flow.async_configure(
result["flow_id"],
{"timeout": "6", "email": "test-email", "password": "test-password"},
)
assert result2["type"] == "form"
assert result2["step_id"] == "validation"
mock_sense.return_value.validate_mfa.side_effect = SenseAuthenticationException
# Try with the WRONG verification code give us the form back again
result3 = await hass.config_entries.flow.async_configure(
result["flow_id"],
{CONF_CODE: "000000"},
)
assert result3["type"] == "form"
assert result3["errors"] == {"base": "invalid_auth"}
assert result3["step_id"] == "validation"
async def test_form_mfa_required_timeout(hass, mock_sense):
"""Test we handle invalid auth."""
result = await hass.config_entries.flow.async_init(
DOMAIN, context={"source": config_entries.SOURCE_USER}
)
mock_sense.return_value.authenticate.side_effect = SenseMFARequiredException
result2 = await hass.config_entries.flow.async_configure(
result["flow_id"],
{"timeout": "6", "email": "test-email", "password": "test-password"},
)
assert result2["type"] == "form"
assert result2["step_id"] == "validation"
mock_sense.return_value.validate_mfa.side_effect = SenseAPITimeoutException
result3 = await hass.config_entries.flow.async_configure(
result["flow_id"],
{CONF_CODE: "000000"},
)
assert result3["type"] == "form"
assert result3["errors"] == {"base": "cannot_connect"}
async def test_form_mfa_required_exception(hass, mock_sense):
"""Test we handle invalid auth."""
result = await hass.config_entries.flow.async_init(
DOMAIN, context={"source": config_entries.SOURCE_USER}
)
mock_sense.return_value.authenticate.side_effect = SenseMFARequiredException
result2 = await hass.config_entries.flow.async_configure(
result["flow_id"],
{"timeout": "6", "email": "test-email", "password": "test-password"},
)
assert result2["type"] == "form"
assert result2["step_id"] == "validation"
mock_sense.return_value.validate_mfa.side_effect = Exception
result3 = await hass.config_entries.flow.async_configure(
result["flow_id"],
{CONF_CODE: "000000"},
)
assert result3["type"] == "form"
assert result3["errors"] == {"base": "unknown"}
async def test_form_cannot_connect(hass): async def test_form_cannot_connect(hass):
"""Test we handle cannot connect error.""" """Test we handle cannot connect error."""
result = await hass.config_entries.flow.async_init( result = await hass.config_entries.flow.async_init(
@ -91,3 +225,57 @@ async def test_form_unknown_exception(hass):
assert result2["type"] == "form" assert result2["type"] == "form"
assert result2["errors"] == {"base": "unknown"} assert result2["errors"] == {"base": "unknown"}
async def test_reauth_no_form(hass, mock_sense):
"""Test reauth where no form needed."""
# set up initially
entry = MockConfigEntry(
domain=DOMAIN,
data=MOCK_CONFIG,
unique_id="test-email",
)
entry.add_to_hass(hass)
with patch(
"homeassistant.config_entries.ConfigEntries.async_reload",
return_value=True,
):
result = await hass.config_entries.flow.async_init(
DOMAIN, context={"source": config_entries.SOURCE_REAUTH}, data=MOCK_CONFIG
)
assert result["type"] == "abort"
assert result["reason"] == "reauth_successful"
async def test_reauth_password(hass, mock_sense):
"""Test reauth form."""
# set up initially
entry = MockConfigEntry(
domain=DOMAIN,
data=MOCK_CONFIG,
unique_id="test-email",
)
entry.add_to_hass(hass)
mock_sense.return_value.authenticate.side_effect = SenseAuthenticationException
# Reauth success without user input
result = await hass.config_entries.flow.async_init(
DOMAIN, context={"source": config_entries.SOURCE_REAUTH}, data=entry.data
)
assert result["type"] == "form"
mock_sense.return_value.authenticate.side_effect = None
with patch(
"homeassistant.components.sense.async_setup_entry",
return_value=True,
):
result2 = await hass.config_entries.flow.async_configure(
result["flow_id"],
{"password": "test-password"},
)
await hass.async_block_till_done()
assert result2["type"] == "abort"
assert result2["reason"] == "reauth_successful"