diff --git a/homeassistant/components/opower/config_flow.py b/homeassistant/components/opower/config_flow.py index fdf007c3b68..9f2ec56423d 100644 --- a/homeassistant/components/opower/config_flow.py +++ b/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( - title=f"{user_input[CONF_UTILITY]} ({user_input[CONF_USERNAME]})", - data=user_input, - ) + return self._async_create_opower_entry(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( - { - vol.Required(CONF_USERNAME): self.reauth_entry.data[CONF_USERNAME], - vol.Required(CONF_PASSWORD): str, - } - ), + data_schema=vol.Schema(schema), errors=errors, ) diff --git a/homeassistant/components/opower/const.py b/homeassistant/components/opower/const.py index b996a214a05..c07d41bbdcf 100644 --- a/homeassistant/components/opower/const.py +++ b/homeassistant/components/opower/const.py @@ -3,3 +3,4 @@ DOMAIN = "opower" CONF_UTILITY = "utility" +CONF_TOTP_SECRET = "totp_secret" diff --git a/homeassistant/components/opower/coordinator.py b/homeassistant/components/opower/coordinator.py index 4e2b68df579..1410b62b7b6 100644 --- a/homeassistant/components/opower/coordinator.py +++ b/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( diff --git a/homeassistant/components/opower/strings.json b/homeassistant/components/opower/strings.json index 037983eb6ff..ac931bf9308 100644 --- a/homeassistant/components/opower/strings.json +++ b/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)" } } }, diff --git a/tests/components/opower/test_config_flow.py b/tests/components/opower/test_config_flow.py index 6a45a0dcc56..0391e42ca16 100644 --- a/tests/components/opower/test_config_flow.py +++ b/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