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",
"name": "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"],
"quality_scale": "internal",
"iot_class": "local_push",

View File

@ -3,18 +3,21 @@ import asyncio
from datetime import timedelta
import logging
from sense_energy import ASyncSenseable, SenseAuthenticationException
from sense_energy import (
ASyncSenseable,
SenseAuthenticationException,
SenseMFARequiredException,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import (
CONF_EMAIL,
CONF_PASSWORD,
CONF_TIMEOUT,
EVENT_HOMEASSISTANT_STOP,
Platform,
)
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.dispatcher import async_dispatcher_send
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
email = entry_data[CONF_EMAIL]
password = entry_data[CONF_PASSWORD]
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)
gateway = ASyncSenseable(
@ -69,16 +75,11 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
gateway.rate_limit = ACTIVE_UPDATE_RATE
try:
await gateway.authenticate(email, password)
except SenseAuthenticationException:
_LOGGER.error("Could not authenticate with sense server")
return False
except SENSE_TIMEOUT_EXCEPTIONS as 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
gateway.load_auth(access_token, user_id, monitor_id)
await gateway.get_monitor_data()
except (SenseAuthenticationException, SenseMFARequiredException) as err:
_LOGGER.warning("Sense authentication expired")
raise ConfigEntryAuthFailed(err) from err
sense_devices_data = SenseDevicesData()
try:
@ -91,11 +92,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
except SENSE_EXCEPTIONS as 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(
hass,
_LOGGER,
name=f"Sense Trends {email}",
update_method=gateway.update_trend_data,
update_method=_async_update_trend,
update_interval=timedelta(seconds=300),
)
# Start out as unavailable so we do not report 0 data

View File

@ -1,11 +1,15 @@
"""Config flow for Sense integration."""
import logging
from sense_energy import ASyncSenseable, SenseAuthenticationException
from sense_energy import (
ASyncSenseable,
SenseAuthenticationException,
SenseMFARequiredException,
)
import voluptuous as vol
from homeassistant import config_entries, core
from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, CONF_TIMEOUT
from homeassistant import config_entries
from homeassistant.const import CONF_CODE, CONF_EMAIL, CONF_PASSWORD, CONF_TIMEOUT
from homeassistant.helpers.aiohttp_client import async_get_clientsession
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):
"""Handle a config flow for Sense."""
VERSION = 1
async def async_step_user(self, user_input=None):
"""Handle the initial step."""
def __init__(self):
"""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 = {}
if user_input is not None:
if user_input:
try:
info = await validate_input(self.hass, user_input)
await self.async_set_unique_id(user_input[CONF_EMAIL])
return self.async_create_entry(title=info["title"], data=user_input)
await self._gateway.validate_mfa(user_input[CONF_CODE])
except SENSE_TIMEOUT_EXCEPTIONS:
errors["base"] = "cannot_connect"
except SenseAuthenticationException:
@ -59,7 +100,43 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
except Exception: # pylint: disable=broad-except
_LOGGER.exception("Unexpected exception")
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(
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",
"name": "Sense",
"documentation": "https://www.home-assistant.io/integrations/sense",
"requirements": ["sense_energy==0.9.6"],
"requirements": ["sense_energy==0.10.2"],
"codeowners": ["@kbickar"],
"config_flow": true,
"dhcp": [

View File

@ -8,6 +8,19 @@
"password": "[%key:common::config_flow::data::password%]",
"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": {
@ -16,7 +29,8 @@
"unknown": "[%key:common::config_flow::error::unknown%]"
},
"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": {
"abort": {
"already_configured": "Device is already configured"
"already_configured": "Device is already configured",
"reauth_successful": "Re-authentication was successful"
},
"error": {
"cannot_connect": "Failed to connect",
@ -9,6 +10,13 @@
"unknown": "Unexpected error"
},
"step": {
"reauth_validate": {
"data": {
"password": "Password"
},
"description": "The Sense integration needs to re-authenticate your account {email}.",
"title": "Reauthenticate Integration"
},
"user": {
"data": {
"email": "Email",
@ -16,6 +24,12 @@
"timeout": "Timeout"
},
"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.sense
sense_energy==0.9.6
sense_energy==0.10.2
# homeassistant.components.sentry
sentry-sdk==1.5.5

View File

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

View File

@ -1,13 +1,44 @@
"""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.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."""
result = await hass.config_entries.flow.async_init(
@ -16,7 +47,7 @@ async def test_form(hass):
assert result["type"] == "form"
assert result["errors"] == {}
with patch("sense_energy.ASyncSenseable.authenticate", return_value=True,), patch(
with patch(
"homeassistant.components.sense.async_setup_entry",
return_value=True,
) as mock_setup_entry:
@ -28,11 +59,7 @@ async def test_form(hass):
assert result2["type"] == "create_entry"
assert result2["title"] == "test-email"
assert result2["data"] == {
"timeout": 6,
"email": "test-email",
"password": "test-password",
}
assert result2["data"] == MOCK_CONFIG
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"}
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):
"""Test we handle cannot connect error."""
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["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"