From da4404e8cf92267c6899a2241c01273a2c32b0e9 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Wed, 20 Jan 2021 22:10:40 +0100 Subject: [PATCH] Introduce reauthentication flow to UniFi integration (#45360) * Improve site selection * Reauth flow and tests Add **kwargs to mock_aiohttp_client create_session to support inputting verify_ssl and cookie_jar * Update homeassistant/components/unifi/config_flow.py Co-authored-by: J. Nick Koston * Minor improvements * Improve coverage Co-authored-by: J. Nick Koston --- homeassistant/components/unifi/config_flow.py | 82 +++++++++++++------ homeassistant/components/unifi/controller.py | 31 +++++-- homeassistant/components/unifi/strings.json | 4 +- .../components/unifi/translations/en.json | 4 +- tests/components/unifi/test_config_flow.py | 53 +++++++++++- tests/components/unifi/test_controller.py | 19 +++++ tests/test_util/aiohttp.py | 2 +- 7 files changed, 157 insertions(+), 38 deletions(-) diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index f29c3951869..6d266025946 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/homeassistant/components/unifi/config_flow.py @@ -66,8 +66,10 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN): def __init__(self): """Initialize the UniFi flow.""" self.config = None - self.desc = None self.sites = None + self.reauth_config_entry = {} + self.reauth_config = {} + self.reauth_schema = {} async def async_step_user(self, user_input=None): """Handle a flow initialized by the user.""" @@ -87,7 +89,13 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN): controller = await get_controller(self.hass, **self.config) - self.sites = await controller.sites() + sites = await controller.sites() + self.sites = {site["name"]: site["desc"] for site in sites.values()} + + if self.reauth_config.get(CONF_SITE_ID) in self.sites: + return await self.async_step_site( + {CONF_SITE_ID: self.reauth_config[CONF_SITE_ID]} + ) return await self.async_step_site() @@ -108,17 +116,17 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN): if await async_discover_unifi(self.hass): host = "unifi" + data = self.reauth_schema or { + vol.Required(CONF_HOST, default=host): str, + vol.Required(CONF_USERNAME): str, + vol.Required(CONF_PASSWORD): str, + vol.Optional(CONF_PORT, default=DEFAULT_PORT): int, + vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): bool, + } + return self.async_show_form( step_id="user", - data_schema=vol.Schema( - { - vol.Required(CONF_HOST, default=host): str, - vol.Required(CONF_USERNAME): str, - vol.Required(CONF_PASSWORD): str, - vol.Optional(CONF_PORT, default=DEFAULT_PORT): int, - vol.Optional(CONF_VERIFY_SSL, default=DEFAULT_VERIFY_SSL): bool, - } - ), + data_schema=vol.Schema(data), errors=errors, ) @@ -128,12 +136,17 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN): if user_input is not None: try: - desc = user_input.get(CONF_SITE_ID, self.desc) + self.config[CONF_SITE_ID] = user_input[CONF_SITE_ID] + data = {CONF_CONTROLLER: self.config} - for site in self.sites.values(): - if desc == site["desc"]: - self.config[CONF_SITE_ID] = site["name"] - break + if self.reauth_config_entry: + self.hass.config_entries.async_update_entry( + self.reauth_config_entry, data=data + ) + await self.hass.config_entries.async_reload( + self.reauth_config_entry.entry_id + ) + return self.async_abort(reason="reauth_successful") for entry in self._async_current_entries(): controller = entry.data[CONF_CONTROLLER] @@ -143,27 +156,44 @@ class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN): ): raise AlreadyConfigured - data = {CONF_CONTROLLER: self.config} - - return self.async_create_entry(title=desc, data=data) + site_nice_name = self.sites[self.config[CONF_SITE_ID]] + return self.async_create_entry(title=site_nice_name, data=data) except AlreadyConfigured: return self.async_abort(reason="already_configured") if len(self.sites) == 1: - self.desc = next(iter(self.sites.values()))["desc"] - return await self.async_step_site(user_input={}) - - sites = [] - for site in self.sites.values(): - sites.append(site["desc"]) + return await self.async_step_site({CONF_SITE_ID: next(iter(self.sites))}) return self.async_show_form( step_id="site", - data_schema=vol.Schema({vol.Required(CONF_SITE_ID): vol.In(sites)}), + data_schema=vol.Schema({vol.Required(CONF_SITE_ID): vol.In(self.sites)}), errors=errors, ) + async def async_step_reauth(self, config_entry: dict): + """Trigger a reauthentication flow.""" + self.reauth_config_entry = config_entry + self.reauth_config = config_entry.data[CONF_CONTROLLER] + + # pylint: disable=no-member # https://github.com/PyCQA/pylint/issues/3167 + self.context["title_placeholders"] = { + CONF_HOST: self.reauth_config[CONF_HOST], + CONF_SITE_ID: config_entry.title, + } + + self.reauth_schema = { + vol.Required(CONF_HOST, default=self.reauth_config[CONF_HOST]): str, + vol.Required(CONF_USERNAME, default=self.reauth_config[CONF_USERNAME]): str, + vol.Required(CONF_PASSWORD): str, + vol.Required(CONF_PORT, default=self.reauth_config[CONF_PORT]): int, + vol.Required( + CONF_VERIFY_SSL, default=self.reauth_config[CONF_VERIFY_SSL] + ): bool, + } + + return await self.async_step_user() + class UnifiOptionsFlowHandler(config_entries.OptionsFlow): """Handle Unifi options.""" diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 095eda9dff9..991940ff721 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -28,6 +28,7 @@ import async_timeout from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN +from homeassistant.config_entries import SOURCE_REAUTH from homeassistant.const import CONF_HOST from homeassistant.core import callback from homeassistant.exceptions import ConfigEntryNotReady @@ -343,8 +344,14 @@ class UniFiController: except CannotConnect as err: raise ConfigEntryNotReady from err - except Exception as err: # pylint: disable=broad-except - LOGGER.error("Unknown error connecting with UniFi controller: %s", err) + except AuthenticationRequired: + self.hass.async_create_task( + self.hass.config_entries.flow.async_init( + UNIFI_DOMAIN, + context={"source": SOURCE_REAUTH}, + data=self.config_entry, + ) + ) return False # Restore clients that is not a part of active clients list. @@ -419,7 +426,13 @@ class UniFiController: @staticmethod async def async_config_entry_updated(hass, config_entry) -> None: - """Handle signals of config entry being updated.""" + """Handle signals of config entry being updated. + + If config entry is updated due to reauth flow + the entry might already have been reset and thus is not available. + """ + if config_entry.entry_id not in hass.data[UNIFI_DOMAIN]: + return controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id] async_dispatcher_send(hass, controller.signal_options_update) @@ -510,7 +523,7 @@ async def get_controller( return controller except aiounifi.Unauthorized as err: - LOGGER.warning("Connected to UniFi at %s but not registered.", host) + LOGGER.warning("Connected to UniFi at %s but not registered: %s", host, err) raise AuthenticationRequired from err except ( @@ -519,9 +532,13 @@ async def get_controller( aiounifi.ServiceUnavailable, aiounifi.RequestError, ) as err: - LOGGER.error("Error connecting to the UniFi controller at %s", host) + LOGGER.error("Error connecting to the UniFi controller at %s: %s", host, err) raise CannotConnect from err - except aiounifi.AiounifiException as err: - LOGGER.exception("Unknown UniFi communication error occurred") + except aiounifi.LoginRequired as err: + LOGGER.warning("Connected to UniFi at %s but login required: %s", host, err) + raise AuthenticationRequired from err + + except aiounifi.AiounifiException as err: + LOGGER.exception("Unknown UniFi communication error occurred: %s", err) raise AuthenticationRequired from err diff --git a/homeassistant/components/unifi/strings.json b/homeassistant/components/unifi/strings.json index 75e23ae2ed1..1bfba247ddc 100644 --- a/homeassistant/components/unifi/strings.json +++ b/homeassistant/components/unifi/strings.json @@ -1,5 +1,6 @@ { "config": { + "flow_title": "{site} ({host})", "step": { "user": { "title": "Set up UniFi Controller", @@ -19,7 +20,8 @@ "unknown_client_mac": "No client available on that MAC address" }, "abort": { - "already_configured": "Controller site is already configured" + "already_configured": "Controller site is already configured", + "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]" } }, "options": { diff --git a/homeassistant/components/unifi/translations/en.json b/homeassistant/components/unifi/translations/en.json index 968d90e377c..d5da167c1b1 100644 --- a/homeassistant/components/unifi/translations/en.json +++ b/homeassistant/components/unifi/translations/en.json @@ -1,13 +1,15 @@ { "config": { "abort": { - "already_configured": "Controller site is already configured" + "already_configured": "Controller site is already configured", + "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]" }, "error": { "faulty_credentials": "Invalid authentication", "service_unavailable": "Failed to connect", "unknown_client_mac": "No client available on that MAC address" }, + "flow_title": "{site} ({host})", "step": { "user": { "data": { diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py index c57024dc758..0626d8d7b86 100644 --- a/tests/components/unifi/test_config_flow.py +++ b/tests/components/unifi/test_config_flow.py @@ -20,6 +20,7 @@ from homeassistant.components.unifi.const import ( CONF_TRACK_WIRED_CLIENTS, DOMAIN as UNIFI_DOMAIN, ) +from homeassistant.config_entries import SOURCE_REAUTH, SOURCE_USER from homeassistant.const import ( CONF_HOST, CONF_PASSWORD, @@ -184,8 +185,8 @@ async def test_flow_works_multiple_sites(hass, aioclient_mock): assert result["type"] == data_entry_flow.RESULT_TYPE_FORM assert result["step_id"] == "site" - assert result["data_schema"]({"site": "site name"}) - assert result["data_schema"]({"site": "site2 name"}) + assert result["data_schema"]({"site": "default"}) + assert result["data_schema"]({"site": "site2"}) async def test_flow_fails_site_already_configured(hass, aioclient_mock): @@ -314,6 +315,54 @@ async def test_flow_fails_unknown_problem(hass, aioclient_mock): assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT +async def test_reauth_flow_update_configuration(hass, aioclient_mock): + """Verify reauth flow can update controller configuration.""" + controller = await setup_unifi_integration(hass) + + result = await hass.config_entries.flow.async_init( + UNIFI_DOMAIN, + context={"source": SOURCE_REAUTH}, + data=controller.config_entry, + ) + + assert result["type"] == data_entry_flow.RESULT_TYPE_FORM + assert result["step_id"] == SOURCE_USER + + aioclient_mock.get("https://1.2.3.4:1234", status=302) + + aioclient_mock.post( + "https://1.2.3.4:1234/api/login", + json={"data": "login successful", "meta": {"rc": "ok"}}, + headers={"content-type": CONTENT_TYPE_JSON}, + ) + + aioclient_mock.get( + "https://1.2.3.4:1234/api/self/sites", + json={ + "data": [{"desc": "Site name", "name": "site_id", "role": "admin"}], + "meta": {"rc": "ok"}, + }, + headers={"content-type": CONTENT_TYPE_JSON}, + ) + + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + user_input={ + CONF_HOST: "1.2.3.4", + CONF_USERNAME: "new_name", + CONF_PASSWORD: "new_pass", + CONF_PORT: 1234, + CONF_VERIFY_SSL: True, + }, + ) + + assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT + assert result["reason"] == "reauth_successful" + assert controller.host == "1.2.3.4" + assert controller.config_entry.data[CONF_CONTROLLER][CONF_USERNAME] == "new_name" + assert controller.config_entry.data[CONF_CONTROLLER][CONF_PASSWORD] == "new_pass" + + async def test_advanced_option_flow(hass): """Test advanced config flow options.""" controller = await setup_unifi_integration( diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index 8105b88a6bf..78741628063 100644 --- a/tests/components/unifi/test_controller.py +++ b/tests/components/unifi/test_controller.py @@ -222,6 +222,17 @@ async def test_controller_not_accessible(hass): assert hass.data[UNIFI_DOMAIN] == {} +async def test_controller_trigger_reauth_flow(hass): + """Failed authentication trigger a reauthentication flow.""" + with patch( + "homeassistant.components.unifi.controller.get_controller", + side_effect=AuthenticationRequired, + ), patch.object(hass.config_entries.flow, "async_init") as mock_flow_init: + await setup_unifi_integration(hass) + mock_flow_init.assert_called_once() + assert hass.data[UNIFI_DOMAIN] == {} + + async def test_controller_unknown_error(hass): """Unknown errors are handled.""" with patch( @@ -319,6 +330,14 @@ async def test_get_controller_controller_unavailable(hass): await get_controller(hass, **CONTROLLER_DATA) +async def test_get_controller_login_required(hass): + """Check that get_controller can handle unknown errors.""" + with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch( + "aiounifi.Controller.login", side_effect=aiounifi.LoginRequired + ), pytest.raises(AuthenticationRequired): + await get_controller(hass, **CONTROLLER_DATA) + + async def test_get_controller_unknown_error(hass): """Check that get_controller can handle unknown errors.""" with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch( diff --git a/tests/test_util/aiohttp.py b/tests/test_util/aiohttp.py index fde1cb6ca2e..53949c20b06 100644 --- a/tests/test_util/aiohttp.py +++ b/tests/test_util/aiohttp.py @@ -276,7 +276,7 @@ def mock_aiohttp_client(): """Context manager to mock aiohttp client.""" mocker = AiohttpClientMocker() - def create_session(hass, *args): + def create_session(hass, *args, **kwargs): session = mocker.create_session(hass.loop) async def close_session(event):