Add support for MFA auth in opower (#97878)

* Add support for MFA auth in opower

* Make MFA an extra step

---------

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

homeassistant/components/opower/config_flow.py

@@ -10,17 +10,19 @@ import voluptuous as vol
 
 from homeassistant import config_entries
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.data_entry_flow import FlowResult
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
 
-from .const import CONF_UTILITY, DOMAIN
+from .const import CONF_TOTP_SECRET, CONF_UTILITY, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
 STEP_USER_DATA_SCHEMA = vol.Schema(
     {
-        vol.Required(CONF_UTILITY): vol.In(get_supported_utility_names()),
+        vol.Required(CONF_UTILITY): vol.In(
+            get_supported_utility_names(supports_mfa=True)
+        ),
         vol.Required(CONF_USERNAME): str,
         vol.Required(CONF_PASSWORD): str,
     }
@@ -36,6 +38,7 @@ async def _validate_login(
         login_data[CONF_UTILITY],
         login_data[CONF_USERNAME],
         login_data[CONF_PASSWORD],
+        login_data.get(CONF_TOTP_SECRET, None),
     )
     errors: dict[str, str] = {}
     try:
@@ -47,6 +50,12 @@ async def _validate_login(
     return errors
 
 
+@callback
+def _supports_mfa(utility: str) -> bool:
+    """Return whether the utility supports MFA."""
+    return utility not in get_supported_utility_names(supports_mfa=False)
+
+
 class OpowerConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a config flow for Opower."""
 
@@ -55,6 +64,7 @@ class OpowerConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     def __init__(self) -> None:
         """Initialize a new OpowerConfigFlow."""
         self.reauth_entry: config_entries.ConfigEntry | None = None
+        self.utility_info: dict[str, Any] | None = None
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
@@ -68,16 +78,56 @@ class OpowerConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                     CONF_USERNAME: user_input[CONF_USERNAME],
                 }
             )
+            if _supports_mfa(user_input[CONF_UTILITY]):
+                self.utility_info = user_input
+                return await self.async_step_mfa()
+
             errors = await _validate_login(self.hass, user_input)
             if not errors:
-                return self.async_create_entry(
+                return self._async_create_opower_entry(user_input)
-                    title=f"{user_input[CONF_UTILITY]} ({user_input[CONF_USERNAME]})",
+
-                    data=user_input,
-                )
         return self.async_show_form(
             step_id="user", data_schema=STEP_USER_DATA_SCHEMA, errors=errors
         )
 
+    async def async_step_mfa(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle MFA step."""
+        assert self.utility_info is not None
+        errors: dict[str, str] = {}
+        if user_input is not None:
+            data = {**self.utility_info, **user_input}
+            errors = await _validate_login(self.hass, data)
+            if not errors:
+                return self._async_create_opower_entry(data)
+
+        if errors:
+            schema = {
+                vol.Required(
+                    CONF_USERNAME, default=self.utility_info[CONF_USERNAME]
+                ): str,
+                vol.Required(CONF_PASSWORD): str,
+            }
+        else:
+            schema = {}
+
+        schema[vol.Required(CONF_TOTP_SECRET)] = str
+
+        return self.async_show_form(
+            step_id="mfa",
+            data_schema=vol.Schema(schema),
+            errors=errors,
+        )
+
+    @callback
+    def _async_create_opower_entry(self, data: dict[str, Any]) -> FlowResult:
+        """Create the config entry."""
+        return self.async_create_entry(
+            title=f"{data[CONF_UTILITY]} ({data[CONF_USERNAME]})",
+            data=data,
+        )
+
     async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
         """Handle configuration by re-auth."""
         self.reauth_entry = self.hass.config_entries.async_get_entry(
@@ -100,13 +150,14 @@ class OpowerConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
                 )
                 await self.hass.config_entries.async_reload(self.reauth_entry.entry_id)
                 return self.async_abort(reason="reauth_successful")
+        schema = {
+            vol.Required(CONF_USERNAME): self.reauth_entry.data[CONF_USERNAME],
+            vol.Required(CONF_PASSWORD): str,
+        }
+        if _supports_mfa(self.reauth_entry.data[CONF_UTILITY]):
+            schema[vol.Optional(CONF_TOTP_SECRET)] = str
         return self.async_show_form(
             step_id="reauth_confirm",
-            data_schema=vol.Schema(
+            data_schema=vol.Schema(schema),
-                {
-                    vol.Required(CONF_USERNAME): self.reauth_entry.data[CONF_USERNAME],
-                    vol.Required(CONF_PASSWORD): str,
-                }
-            ),
             errors=errors,
         )

homeassistant/components/opower/const.py

@@ -3,3 +3,4 @@
 DOMAIN = "opower"
 
 CONF_UTILITY = "utility"
+CONF_TOTP_SECRET = "totp_secret"

homeassistant/components/opower/coordinator.py

@@ -28,7 +28,7 @@ from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers import aiohttp_client
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import CONF_UTILITY, DOMAIN
+from .const import CONF_TOTP_SECRET, CONF_UTILITY, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -55,6 +55,7 @@ class OpowerCoordinator(DataUpdateCoordinator[dict[str, Forecast]]):
             entry_data[CONF_UTILITY],
             entry_data[CONF_USERNAME],
             entry_data[CONF_PASSWORD],
+            entry_data.get(CONF_TOTP_SECRET, None),
         )
 
     async def _async_update_data(

homeassistant/components/opower/strings.json

@@ -5,14 +5,16 @@
         "data": {
           "utility": "Utility name",
           "username": "[%key:common::config_flow::data::username%]",
-          "password": "[%key:common::config_flow::data::password%]"
+          "password": "[%key:common::config_flow::data::password%]",
+          "totp_secret": "TOTP Secret (only for some utilities, see documentation)"
         }
       },
       "reauth_confirm": {
         "title": "[%key:common::config_flow::title::reauth%]",
         "data": {
           "username": "[%key:common::config_flow::data::username%]",
-          "password": "[%key:common::config_flow::data::password%]"
+          "password": "[%key:common::config_flow::data::password%]",
+          "totp_secret": "TOTP Secret (only for some utilities, see documentation)"
         }
       }
     },

tests/components/opower/test_config_flow.py

@@ -68,6 +68,110 @@ async def test_form(
     assert mock_login.call_count == 1
 
 
+async def test_form_with_mfa(
+    recorder_mock: Recorder, hass: HomeAssistant, mock_setup_entry: AsyncMock
+) -> None:
+    """Test we get the form."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == FlowResultType.FORM
+    assert not result["errors"]
+
+    result2 = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {
+            "utility": "Consolidated Edison (ConEd)",
+            "username": "test-username",
+            "password": "test-password",
+        },
+    )
+    assert result2["type"] == FlowResultType.FORM
+    assert not result2["errors"]
+
+    with patch(
+        "homeassistant.components.opower.config_flow.Opower.async_login",
+    ) as mock_login:
+        result3 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                "totp_secret": "test-totp",
+            },
+        )
+
+    assert result3["type"] == FlowResultType.CREATE_ENTRY
+    assert result3["title"] == "Consolidated Edison (ConEd) (test-username)"
+    assert result3["data"] == {
+        "utility": "Consolidated Edison (ConEd)",
+        "username": "test-username",
+        "password": "test-password",
+        "totp_secret": "test-totp",
+    }
+    assert len(mock_setup_entry.mock_calls) == 1
+    assert mock_login.call_count == 1
+
+
+async def test_form_with_mfa_bad_secret(
+    recorder_mock: Recorder, hass: HomeAssistant, mock_setup_entry: AsyncMock
+) -> None:
+    """Test MFA asks for password again when validation fails."""
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == FlowResultType.FORM
+    assert not result["errors"]
+
+    result2 = await hass.config_entries.flow.async_configure(
+        result["flow_id"],
+        {
+            "utility": "Consolidated Edison (ConEd)",
+            "username": "test-username",
+            "password": "test-password",
+        },
+    )
+    assert result2["type"] == FlowResultType.FORM
+    assert not result2["errors"]
+
+    with patch(
+        "homeassistant.components.opower.config_flow.Opower.async_login",
+        side_effect=InvalidAuth,
+    ) as mock_login:
+        result3 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                "totp_secret": "test-totp",
+            },
+        )
+
+    assert result3["type"] == FlowResultType.FORM
+    assert result3["errors"] == {
+        "base": "invalid_auth",
+    }
+
+    with patch(
+        "homeassistant.components.opower.config_flow.Opower.async_login",
+    ) as mock_login:
+        result4 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                "username": "test-username",
+                "password": "updated-password",
+                "totp_secret": "updated-totp",
+            },
+        )
+
+    assert result4["type"] == FlowResultType.CREATE_ENTRY
+    assert result4["title"] == "Consolidated Edison (ConEd) (test-username)"
+    assert result4["data"] == {
+        "utility": "Consolidated Edison (ConEd)",
+        "username": "test-username",
+        "password": "updated-password",
+        "totp_secret": "updated-totp",
+    }
+    assert len(mock_setup_entry.mock_calls) == 1
+    assert mock_login.call_count == 1
+
+
 @pytest.mark.parametrize(
     ("api_exception", "expected_error"),
     [
@@ -204,3 +308,54 @@ async def test_form_valid_reauth(
     assert len(mock_unload_entry.mock_calls) == 1
     assert len(mock_setup_entry.mock_calls) == 1
     assert mock_login.call_count == 1
+
+
+async def test_form_valid_reauth_with_mfa(
+    recorder_mock: Recorder,
+    hass: HomeAssistant,
+    mock_setup_entry: AsyncMock,
+    mock_unload_entry: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+) -> None:
+    """Test that we can handle a valid reauth."""
+    hass.config_entries.async_update_entry(
+        mock_config_entry,
+        data={
+            **mock_config_entry.data,
+            # Requires MFA
+            "utility": "Consolidated Edison (ConEd)",
+        },
+    )
+    mock_config_entry.state = ConfigEntryState.LOADED
+    mock_config_entry.async_start_reauth(hass)
+    await hass.async_block_till_done()
+
+    flows = hass.config_entries.flow.async_progress()
+    assert len(flows) == 1
+    result = flows[0]
+
+    with patch(
+        "homeassistant.components.opower.config_flow.Opower.async_login",
+    ) as mock_login:
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                "username": "test-username",
+                "password": "test-password2",
+                "totp_secret": "test-totp",
+            },
+        )
+
+    assert result["type"] == FlowResultType.ABORT
+    assert result["reason"] == "reauth_successful"
+
+    await hass.async_block_till_done()
+    assert hass.config_entries.async_entries(DOMAIN)[0].data == {
+        "utility": "Consolidated Edison (ConEd)",
+        "username": "test-username",
+        "password": "test-password2",
+        "totp_secret": "test-totp",
+    }
+    assert len(mock_unload_entry.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1
+    assert mock_login.call_count == 1