Update Notion auth to store refresh tokens instead of account passwords (#109670)

.coveragerc

@@ -873,6 +873,7 @@ omit =
     homeassistant/components/notion/__init__.py
     homeassistant/components/notion/binary_sensor.py
     homeassistant/components/notion/sensor.py
+    homeassistant/components/notion/util.py
     homeassistant/components/nsw_fuel_station/sensor.py
     homeassistant/components/nuki/__init__.py
     homeassistant/components/nuki/binary_sensor.py

homeassistant/components/notion/__init__.py

@@ -7,7 +7,6 @@ from datetime import timedelta
 from typing import Any
 from uuid import UUID
 
-from aionotion import async_get_client
 from aionotion.bridge.models import Bridge
 from aionotion.errors import InvalidCredentialsError, NotionError
 from aionotion.listener.models import Listener, ListenerKind
@@ -19,7 +18,6 @@ from homeassistant.const import CONF_PASSWORD, CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers import (
-    aiohttp_client,
     config_validation as cv,
     device_registry as dr,
     entity_registry as er,
@@ -33,6 +31,8 @@ from homeassistant.helpers.update_coordinator import (
 )
 
 from .const import (
+    CONF_REFRESH_TOKEN,
+    CONF_USER_UUID,
     DOMAIN,
     LOGGER,
     SENSOR_BATTERY,
@@ -46,6 +46,7 @@ from .const import (
     SENSOR_TEMPERATURE,
     SENSOR_WINDOW_HINGED,
 )
+from .util import async_get_client_with_credentials, async_get_client_with_refresh_token
 
 PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR]
 
@@ -139,25 +140,48 @@ class NotionData:
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Notion as a config entry."""
-    if not entry.unique_id:
+    entry_updates: dict[str, Any] = {"data": {**entry.data}}
-        hass.config_entries.async_update_entry(
-            entry, unique_id=entry.data[CONF_USERNAME]
-        )
 
-    session = aiohttp_client.async_get_clientsession(hass)
+    if not entry.unique_id:
+        entry_updates["unique_id"] = entry.data[CONF_USERNAME]
 
     try:
-        client = await async_get_client(
+        if password := entry_updates["data"].pop(CONF_PASSWORD, None):
-            entry.data[CONF_USERNAME],
+            # If a password exists in the config entry data, use it to get a new client
-            entry.data[CONF_PASSWORD],
+            # (and pop it from the new entry data):
-            session=session,
+            client = await async_get_client_with_credentials(
-            use_legacy_auth=True,
+                hass, entry.data[CONF_USERNAME], password
-        )
+            )
+        else:
+            # If a password doesn't exist in the config entry data, we can safely assume
+            # that a refresh token and user UUID do, so we use them to get the client:
+            client = await async_get_client_with_refresh_token(
+                hass,
+                entry.data[CONF_USER_UUID],
+                entry.data[CONF_REFRESH_TOKEN],
+            )
     except InvalidCredentialsError as err:
-        raise ConfigEntryAuthFailed("Invalid username and/or password") from err
+        raise ConfigEntryAuthFailed("Invalid credentials") from err
     except NotionError as err:
         raise ConfigEntryNotReady("Config entry failed to load") from err
 
+    # Always update the config entry with the latest refresh token and user UUID:
+    entry_updates["data"][CONF_REFRESH_TOKEN] = client.refresh_token
+    entry_updates["data"][CONF_USER_UUID] = client.user_uuid
+
+    @callback
+    def async_save_refresh_token(refresh_token: str) -> None:
+        """Save a refresh token to the config entry data."""
+        LOGGER.debug("Saving new refresh token to HASS storage")
+        hass.config_entries.async_update_entry(
+            entry, data={**entry.data, CONF_REFRESH_TOKEN: refresh_token}
+        )
+
+    # Create a callback to save the refresh token when it changes:
+    entry.async_on_unload(client.add_refresh_token_callback(async_save_refresh_token))
+
+    hass.config_entries.async_update_entry(entry, **entry_updates)
+
     async def async_update() -> NotionData:
         """Get the latest data from the Notion API."""
         data = NotionData(hass=hass, entry=entry)

homeassistant/components/notion/config_flow.py

@@ -2,9 +2,9 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
+from dataclasses import dataclass, field
 from typing import Any
 
-from aionotion import async_get_client
 from aionotion.errors import InvalidCredentialsError, NotionError
 import voluptuous as vol
 
@@ -13,9 +13,9 @@ from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResult
-from homeassistant.helpers import aiohttp_client
 
-from .const import DOMAIN, LOGGER
+from .const import CONF_REFRESH_TOKEN, CONF_USER_UUID, DOMAIN, LOGGER
+from .util import async_get_client_with_credentials
 
 AUTH_SCHEMA = vol.Schema(
     {
@@ -30,17 +30,23 @@ REAUTH_SCHEMA = vol.Schema(
 )
 
 
+@dataclass(frozen=True, kw_only=True)
+class CredentialsValidationResult:
+    """Define a validation result."""
+
+    user_uuid: str | None = None
+    refresh_token: str | None = None
+    errors: dict[str, Any] = field(default_factory=dict)
+
+
 async def async_validate_credentials(
     hass: HomeAssistant, username: str, password: str
-) -> dict[str, Any]:
+) -> CredentialsValidationResult:
-    """Validate a Notion username and password (returning any errors)."""
+    """Validate a Notion username and password."""
-    session = aiohttp_client.async_get_clientsession(hass)
     errors = {}
 
     try:
-        await async_get_client(
+        client = await async_get_client_with_credentials(hass, username, password)
-            username, password, session=session, use_legacy_auth=True
-        )
     except InvalidCredentialsError:
         errors["base"] = "invalid_auth"
     except NotionError as err:
@@ -50,7 +56,12 @@ async def async_validate_credentials(
         LOGGER.exception("Unknown error while validation credentials: %s", err)
         errors["base"] = "unknown"
 
-    return errors
+    if errors:
+        return CredentialsValidationResult(errors=errors)
+
+    return CredentialsValidationResult(
+        user_uuid=client.user_uuid, refresh_token=client.refresh_token
+    )
 
 
 class NotionFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
@@ -84,20 +95,24 @@ class NotionFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
                 },
             )
 
-        if errors := await async_validate_credentials(
+        credentials_validation_result = await async_validate_credentials(
             self.hass, self._reauth_entry.data[CONF_USERNAME], user_input[CONF_PASSWORD]
-        ):
+        )
+
+        if credentials_validation_result.errors:
             return self.async_show_form(
                 step_id="reauth_confirm",
                 data_schema=REAUTH_SCHEMA,
-                errors=errors,
+                errors=credentials_validation_result.errors,
                 description_placeholders={
                     CONF_USERNAME: self._reauth_entry.data[CONF_USERNAME]
                 },
             )
 
         self.hass.config_entries.async_update_entry(
-            self._reauth_entry, data=self._reauth_entry.data | user_input
+            self._reauth_entry,
+            data=self._reauth_entry.data
+            | {CONF_REFRESH_TOKEN: credentials_validation_result.refresh_token},
         )
         self.hass.async_create_task(
             self.hass.config_entries.async_reload(self._reauth_entry.entry_id)
@@ -114,13 +129,22 @@ class NotionFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
         await self.async_set_unique_id(user_input[CONF_USERNAME])
         self._abort_if_unique_id_configured()
 
-        if errors := await async_validate_credentials(
+        credentials_validation_result = await async_validate_credentials(
             self.hass, user_input[CONF_USERNAME], user_input[CONF_PASSWORD]
-        ):
+        )
+
+        if credentials_validation_result.errors:
             return self.async_show_form(
                 step_id="user",
                 data_schema=AUTH_SCHEMA,
-                errors=errors,
+                errors=credentials_validation_result.errors,
             )
 
-        return self.async_create_entry(title=user_input[CONF_USERNAME], data=user_input)
+        return self.async_create_entry(
+            title=user_input[CONF_USERNAME],
+            data={
+                CONF_USERNAME: user_input[CONF_USERNAME],
+                CONF_USER_UUID: credentials_validation_result.user_uuid,
+                CONF_REFRESH_TOKEN: credentials_validation_result.refresh_token,
+            },
+        )

homeassistant/components/notion/const.py

@@ -4,6 +4,9 @@ import logging
 DOMAIN = "notion"
 LOGGER = logging.getLogger(__package__)
 
+CONF_REFRESH_TOKEN = "refresh_token"
+CONF_USER_UUID = "user_uuid"
+
 SENSOR_BATTERY = "low_battery"
 SENSOR_DOOR = "door"
 SENSOR_GARAGE_DOOR = "garage_door"

homeassistant/components/notion/diagnostics.py

@@ -5,12 +5,12 @@ from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_EMAIL, CONF_PASSWORD, CONF_UNIQUE_ID, CONF_USERNAME
+from homeassistant.const import CONF_EMAIL, CONF_UNIQUE_ID, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
 from . import NotionData
-from .const import DOMAIN
+from .const import CONF_REFRESH_TOKEN, CONF_USER_UUID, DOMAIN
 
 CONF_DEVICE_KEY = "device_key"
 CONF_HARDWARE_ID = "hardware_id"
@@ -23,12 +23,13 @@ TO_REDACT = {
     CONF_EMAIL,
     CONF_HARDWARE_ID,
     CONF_LAST_BRIDGE_HARDWARE_ID,
-    CONF_PASSWORD,
+    CONF_REFRESH_TOKEN,
     # Config entry title and unique ID may contain sensitive data:
     CONF_TITLE,
     CONF_UNIQUE_ID,
     CONF_USERNAME,
     CONF_USER_ID,
+    CONF_USER_UUID,
 }
 
 

homeassistant/components/notion/manifest.json

@@ -7,5 +7,5 @@
   "integration_type": "hub",
   "iot_class": "cloud_polling",
   "loggers": ["aionotion"],
-  "requirements": ["aionotion==2024.02.0"]
+  "requirements": ["aionotion==2024.02.1"]
 }

homeassistant/components/notion/util.py

@@ -0,0 +1,30 @@
+"""Define notion utilities."""
+from aionotion import (
+    async_get_client_with_credentials as cwc,
+    async_get_client_with_refresh_token as cwrt,
+)
+from aionotion.client import Client
+
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import aiohttp_client
+from homeassistant.helpers.instance_id import async_get
+
+
+async def async_get_client_with_credentials(
+    hass: HomeAssistant, email: str, password: str
+) -> Client:
+    """Get a Notion client with credentials."""
+    session = aiohttp_client.async_get_clientsession(hass)
+    instance_id = await async_get(hass)
+    return await cwc(email, password, session=session, session_name=instance_id)
+
+
+async def async_get_client_with_refresh_token(
+    hass: HomeAssistant, user_uuid: str, refresh_token: str
+) -> Client:
+    """Get a Notion client with credentials."""
+    session = aiohttp_client.async_get_clientsession(hass)
+    instance_id = await async_get(hass)
+    return await cwrt(
+        user_uuid, refresh_token, session=session, session_name=instance_id
+    )

requirements_all.txt

@@ -315,7 +315,7 @@ aiomusiccast==0.14.8
 aionanoleaf==0.2.1
 
 # homeassistant.components.notion
-aionotion==2024.02.0
+aionotion==2024.02.1
 
 # homeassistant.components.oncue
 aiooncue==0.3.5

requirements_test_all.txt

@@ -288,7 +288,7 @@ aiomusiccast==0.14.8
 aionanoleaf==0.2.1
 
 # homeassistant.components.notion
-aionotion==2024.02.0
+aionotion==2024.02.1
 
 # homeassistant.components.oncue
 aiooncue==0.3.5

tests/components/notion/conftest.py

@@ -17,6 +17,8 @@ from tests.common import MockConfigEntry, load_fixture
 
 TEST_USERNAME = "user@host.com"
 TEST_PASSWORD = "password123"
+TEST_REFRESH_TOKEN = "abcde12345"
+TEST_USER_UUID = "xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx"
 
 
 @pytest.fixture
@@ -47,6 +49,7 @@ def client_fixture(data_bridge, data_listener, data_sensor, data_user_preference
                 ]
             )
         ),
+        refresh_token=TEST_REFRESH_TOKEN,
         sensor=Mock(
             async_all=AsyncMock(
                 return_value=[
@@ -61,6 +64,7 @@ def client_fixture(data_bridge, data_listener, data_sensor, data_user_preference
                 )
             )
         ),
+        user_uuid=TEST_USER_UUID,
     )
 
 
@@ -107,7 +111,7 @@ def data_user_preferences_fixture():
 
 @pytest.fixture(name="get_client")
 def get_client_fixture(client):
-    """Define a fixture to mock the async_get_client method."""
+    """Define a fixture to mock the client retrieval methods."""
     return AsyncMock(return_value=client)
 
 
@@ -115,10 +119,13 @@ def get_client_fixture(client):
 async def mock_aionotion_fixture(client):
     """Define a fixture to patch aionotion."""
     with patch(
-        "homeassistant.components.notion.async_get_client",
+        "homeassistant.components.notion.async_get_client_with_credentials",
         AsyncMock(return_value=client),
     ), patch(
-        "homeassistant.components.notion.config_flow.async_get_client",
+        "homeassistant.components.notion.async_get_client_with_refresh_token",
+        AsyncMock(return_value=client),
+    ), patch(
+        "homeassistant.components.notion.config_flow.async_get_client_with_credentials",
         AsyncMock(return_value=client),
     ):
         yield

tests/components/notion/test_config_flow.py

@@ -5,12 +5,12 @@ from aionotion.errors import InvalidCredentialsError, NotionError
 import pytest
 
 from homeassistant import data_entry_flow
-from homeassistant.components.notion import DOMAIN
+from homeassistant.components.notion import CONF_REFRESH_TOKEN, CONF_USER_UUID, DOMAIN
 from homeassistant.config_entries import SOURCE_REAUTH, SOURCE_USER
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
 
-from .conftest import TEST_PASSWORD, TEST_USERNAME
+from .conftest import TEST_REFRESH_TOKEN, TEST_USER_UUID, TEST_USERNAME
 
 pytestmark = pytest.mark.usefixtures("mock_setup_entry")
 
@@ -40,7 +40,7 @@ async def test_create_entry(
 
     # Test errors that can arise when getting a Notion API client:
     with patch(
-        "homeassistant.components.notion.config_flow.async_get_client",
+        "homeassistant.components.notion.config_flow.async_get_client_with_credentials",
         get_client_with_exception,
     ):
         result = await hass.config_entries.flow.async_init(
@@ -55,8 +55,9 @@ async def test_create_entry(
     assert result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY
     assert result["title"] == TEST_USERNAME
     assert result["data"] == {
+        CONF_REFRESH_TOKEN: TEST_REFRESH_TOKEN,
         CONF_USERNAME: TEST_USERNAME,
-        CONF_PASSWORD: TEST_PASSWORD,
+        CONF_USER_UUID: TEST_USER_UUID,
     }
 
 
@@ -99,7 +100,7 @@ async def test_reauth(
 
     # Test errors that can arise when getting a Notion API client:
     with patch(
-        "homeassistant.components.notion.config_flow.async_get_client",
+        "homeassistant.components.notion.config_flow.async_get_client_with_credentials",
         get_client_with_exception,
     ):
         result = await hass.config_entries.flow.async_configure(

tests/components/notion/test_diagnostics.py

@@ -21,7 +21,11 @@ async def test_entry_diagnostics(
             "minor_version": 1,
             "domain": DOMAIN,
             "title": REDACTED,
-            "data": {"username": REDACTED, "password": REDACTED},
+            "data": {
+                "refresh_token": REDACTED,
+                "user_uuid": REDACTED,
+                "username": REDACTED,
+            },
             "options": {},
             "pref_disable_new_entities": False,
             "pref_disable_polling": False,