Add Multi factor authentication support for Sense (#66498)

Co-authored-by: J. Nick Koston <nick@koston.org>

homeassistant/components/emulated_kasa/manifest.json

@@ -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",

homeassistant/components/sense/__init__.py

@@ -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)
+        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()
     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

homeassistant/components/sense/config_flow.py

@@ -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 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 .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):
+    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 = {}
-        if user_input is not None:
+        if user_input:
             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:
                 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],
+            },
+        )

homeassistant/components/sense/manifest.json

@@ -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": [

homeassistant/components/sense/strings.json

@@ -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%]"
     }
   }
 }

homeassistant/components/sense/translations/en.json

@@ -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"
             }
         }
     }

requirements_all.txt

@@ -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

requirements_test_all.txt

@@ -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

tests/components/sense/test_config_flow.py

@@ -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"] == {
+    assert result2["data"] == MOCK_CONFIG
-        "timeout": 6,
-        "email": "test-email",
-        "password": "test-password",
-    }
     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"