Replace access token authentication with OAuth2 in Teslemetry

homeassistant/components/teslemetry/__init__.py

@@ -4,6 +4,7 @@
 from collections.abc import Callable
 from typing import Final
 
+from aiohttp import ClientResponseError
 from tesla_fleet_api.const import Scope
 from tesla_fleet_api.exceptions import (
     Forbidden,
@@ -14,16 +15,24 @@
 from tesla_fleet_api.teslemetry import Teslemetry
 from teslemetry_stream import TeslemetryStream
 
+from homeassistant.components.application_credentials import (
+    ClientCredential,
+    async_import_client_credential,
+)
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_ACCESS_TOKEN, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers import config_validation as cv, device_registry as dr
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.config_entry_oauth2_flow import (
+    OAuth2Session,
+    async_get_config_entry_implementation,
+)
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.typing import ConfigType
 
-from .const import DOMAIN, LOGGER
+from .const import CLIENT_ID, DOMAIN, LOGGER
 from .coordinator import (
     TeslemetryEnergyHistoryCoordinator,
     TeslemetryEnergySiteInfoCoordinator,
@@ -56,20 +65,37 @@
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up the Telemetry integration."""
+    await async_import_client_credential(
+        hass,
+        DOMAIN,
+        ClientCredential(CLIENT_ID, "", name="Teslemetry"),
+    )
     async_setup_services(hass)
     return True
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: TeslemetryConfigEntry) -> bool:
     """Set up Teslemetry config."""
 
-    access_token = entry.data[CONF_ACCESS_TOKEN]
     session = async_get_clientsession(hass)
 
+    implementation = await async_get_config_entry_implementation(hass, entry)
+    oauth_session = OAuth2Session(hass, entry, implementation)
+
+    async def _get_access_token() -> str:
+        try:
+            await oauth_session.async_ensure_token_valid()
+        except ClientResponseError as e:
+            if e.status == 401:
+                raise ConfigEntryAuthFailed from e
+            raise ConfigEntryNotReady from e
+        token: str = oauth_session.token[CONF_ACCESS_TOKEN]
+        return token
+
     # Create API connection
     teslemetry = Teslemetry(
         session=session,
-        access_token=access_token,
+        access_token=_get_access_token,
     )
     try:
         calls = await asyncio.gather(
@@ -125,7 +151,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: TeslemetryConfigEntry) -
             if not stream:
                 stream = TeslemetryStream(
                     session,
-                    access_token,
+                    _get_access_token,
                     server=f"{region.lower()}.teslemetry.com",
                     parse_timestamp=True,
                     manual=True,
@@ -276,23 +302,29 @@ async def async_migrate_entry(
     hass: HomeAssistant, config_entry: TeslemetryConfigEntry
 ) -> bool:
     """Migrate config entry."""
-    if config_entry.version > 1:
+    if config_entry.version > 2:
+        # This means the user has downgraded from a future version
         return False
 
-    if config_entry.version == 1 and config_entry.minor_version < 2:
-        # Add unique_id to existing entry
-        teslemetry = Teslemetry(
-            session=async_get_clientsession(hass),
-            access_token=config_entry.data[CONF_ACCESS_TOKEN],
-        )
+    if config_entry.version == 1:
+        access_token = config_entry.data[CONF_ACCESS_TOKEN]
+        session = async_get_clientsession(hass)
+
+        # Convert legacy access token to OAuth tokens using migrate endpoint
         try:
-            metadata = await teslemetry.metadata()
-        except TeslaFleetError as e:
-            LOGGER.error(e.message)
-            return False
+            data = await Teslemetry(session, access_token).migrate_to_oauth(
+                CLIENT_ID, access_token, hass.config.location_name
+            )
+        except ClientResponseError as e:
+            raise ConfigEntryAuthFailed from e
+
+        # Add auth_implementation for OAuth2 flow compatibility
+        data["auth_implementation"] = DOMAIN
 
-        hass.config_entries.async_update_entry(
-            config_entry, unique_id=metadata["uid"], version=1, minor_version=2
+        return hass.config_entries.async_update_entry(
+            config_entry,
+            data=data,
+            version=2,
         )
     return True
 

homeassistant/components/teslemetry/application_credentials.py

@@ -0,0 +1,30 @@
+"""Application Credentials platform the Teslemetry integration."""
+
+from homeassistant.components.application_credentials import (
+    AuthorizationServer,
+    ClientCredential,
+)
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import config_entry_oauth2_flow
+
+from .const import AUTHORIZE_URL, TOKEN_URL
+from .oauth import TeslemetryImplementation
+
+
+async def async_get_authorization_server(hass: HomeAssistant) -> AuthorizationServer:
+    """Return authorization server."""
+    return AuthorizationServer(
+        authorize_url=AUTHORIZE_URL,
+        token_url=TOKEN_URL,
+    )
+
+
+async def async_get_auth_implementation(
+    hass: HomeAssistant, auth_domain: str, credential: ClientCredential
+) -> config_entry_oauth2_flow.AbstractOAuth2Implementation:
+    """Return auth implementation."""
+    return TeslemetryImplementation(
+        hass,
+        auth_domain,
+        credential.client_id,
+    )

homeassistant/components/teslemetry/config_flow.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
+import logging
 from typing import Any
 
 from aiohttp import ClientConnectionError
@@ -12,90 +13,111 @@
     TeslaFleetError,
 )
 from tesla_fleet_api.teslemetry import Teslemetry
-import voluptuous as vol
 
-from homeassistant.config_entries import ConfigFlow, ConfigFlowResult
-from homeassistant.const import CONF_ACCESS_TOKEN
+from homeassistant.components.application_credentials import (
+    ClientCredential,
+    async_import_client_credential,
+)
+from homeassistant.config_entries import SOURCE_REAUTH, ConfigFlowResult
+from homeassistant.helpers import config_entry_oauth2_flow
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .const import DOMAIN, LOGGER
+from .const import CLIENT_ID, DOMAIN, LOGGER
+
+
+class OAuth2FlowHandler(
+    config_entry_oauth2_flow.AbstractOAuth2FlowHandler, domain=DOMAIN
+):
+    """Config flow to handle Teslemetry OAuth2 authentication."""
+
+    DOMAIN = DOMAIN
+    VERSION = 2
 
-TESLEMETRY_SCHEMA = vol.Schema({vol.Required(CONF_ACCESS_TOKEN): str})
-DESCRIPTION_PLACEHOLDERS = {
-    "name": "Teslemetry",
-    "short_url": "teslemetry.com/console",
-    "url": "[teslemetry.com/console](https://teslemetry.com/console)",
-}
+    def __init__(self) -> None:
+        """Initialize config flow."""
+        super().__init__()
+        self.data: dict[str, Any] = {}
+        self.uid: str | None = None
 
+    @property
+    def logger(self) -> logging.Logger:
+        """Return logger."""
+        return LOGGER
+
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
+        """Handle a flow start."""
+        await async_import_client_credential(
+            self.hass,
+            DOMAIN,
+            ClientCredential(CLIENT_ID, "", name="Teslemetry"),
+        )
+        return await super().async_step_user()
+
+    async def async_oauth_create_entry(
+        self,
+        data: dict[str, Any],
+    ) -> ConfigFlowResult:
+        """Handle OAuth completion and create config entry."""
+        self.data = data
+
+        # Test the connection with the OAuth token
+        errors = await self.async_test_connection(data)
+        if errors:
+            return self.async_abort(reason="oauth_error")
+
+        await self.async_set_unique_id(self.uid)
+        if self.source == SOURCE_REAUTH:
+            self._abort_if_unique_id_mismatch(reason="reauth_account_mismatch")
+            return self.async_update_reload_and_abort(
+                self._get_reauth_entry(), data=data
+            )
+        self._abort_if_unique_id_configured()
 
-class TeslemetryConfigFlow(ConfigFlow, domain=DOMAIN):
-    """Config Teslemetry API connection."""
+        return self.async_create_entry(
+            title="Teslemetry",
+            data=data,
+        )
 
-    VERSION = 1
-    MINOR_VERSION = 2
+    async def async_test_connection(self, token_data: dict[str, Any]) -> dict[str, str]:
+        """Test the connection with OAuth token."""
+        access_token = token_data["token"]["access_token"]
 
-    async def async_auth(self, user_input: Mapping[str, Any]) -> dict[str, str]:
-        """Reusable Auth Helper."""
         teslemetry = Teslemetry(
             session=async_get_clientsession(self.hass),
-            access_token=user_input[CONF_ACCESS_TOKEN],
+            access_token=access_token,
         )
+
         try:
             metadata = await teslemetry.metadata()
         except InvalidToken:
-            return {CONF_ACCESS_TOKEN: "invalid_access_token"}
+            return {"base": "invalid_access_token"}
         except SubscriptionRequired:
             return {"base": "subscription_required"}
         except ClientConnectionError:
             return {"base": "cannot_connect"}
         except TeslaFleetError as e:
-            LOGGER.error(e)
+            LOGGER.error("Teslemetry API error: %s", e)
             return {"base": "unknown"}
 
-        await self.async_set_unique_id(metadata["uid"])
+        self.uid = metadata["uid"]
         return {}
 
-    async def async_step_user(
-        self, user_input: Mapping[str, Any] | None = None
-    ) -> ConfigFlowResult:
-        """Get configuration from the user."""
-        errors: dict[str, str] = {}
-        if user_input and not (errors := await self.async_auth(user_input)):
-            self._abort_if_unique_id_configured()
-            return self.async_create_entry(
-                title="Teslemetry",
-                data=user_input,
-            )
-
-        return self.async_show_form(
-            step_id="user",
-            data_schema=TESLEMETRY_SCHEMA,
-            description_placeholders=DESCRIPTION_PLACEHOLDERS,
-            errors=errors,
-        )
-
     async def async_step_reauth(
         self, entry_data: Mapping[str, Any]
     ) -> ConfigFlowResult:
         """Handle reauth on failure."""
         return await self.async_step_reauth_confirm()
 
     async def async_step_reauth_confirm(
-        self, user_input: Mapping[str, Any] | None = None
+        self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
-        """Handle users reauth credentials."""
-
-        errors: dict[str, str] = {}
-
-        if user_input and not (errors := await self.async_auth(user_input)):
-            return self.async_update_reload_and_abort(
-                self._get_reauth_entry(),
-                data=user_input,
+        """Confirm reauth dialog."""
+        if user_input is None:
+            return self.async_show_form(
+                step_id="reauth_confirm",
+                description_placeholders={"name": "Teslemetry"},
             )
 
-        return self.async_show_form(
-            step_id="reauth_confirm",
-            description_placeholders=DESCRIPTION_PLACEHOLDERS,
-            data_schema=TESLEMETRY_SCHEMA,
-            errors=errors,
-        )
+        return await super().async_step_user()

homeassistant/components/teslemetry/const.py

@@ -9,6 +9,11 @@
 
 LOGGER = logging.getLogger(__package__)
 
+# OAuth
+AUTHORIZE_URL = "https://teslemetry.com/connect"
+TOKEN_URL = "https://api.teslemetry.com/oauth/token"
+CLIENT_ID = "homeassistant"
+
 ENERGY_HISTORY_FIELDS = [
     "solar_energy_exported",
     "generator_energy_exported",

homeassistant/components/teslemetry/manifest.json

@@ -3,6 +3,7 @@
   "name": "Teslemetry",
   "codeowners": ["@Bre77"],
   "config_flow": true,
+  "dependencies": ["application_credentials"],
   "documentation": "https://www.home-assistant.io/integrations/teslemetry",
   "integration_type": "hub",
   "iot_class": "cloud_polling",