Feat/auto token renewal

custom_components/iec/config_flow.py

@@ -2,8 +2,9 @@
 
 from __future__ import annotations
 
-import logging
 import asyncio
+import json
+import logging
 from collections.abc import Mapping
 from typing import Any
 
@@ -21,6 +22,9 @@
 from .const import (
     CONF_AVAILABLE_CONTRACTS,
     CONF_BP_NUMBER,
+    CONF_GMAIL_AUTO_REAUTH,
+    CONF_GMAIL_CREDENTIALS,
+    CONF_GMAIL_TOKEN,
     CONF_SELECTED_CONTRACTS,
     CONF_TOTP_SECRET,
     CONF_USER_ID,
@@ -85,6 +89,14 @@ def __init__(self) -> None:
         self.data: dict[str, Any] | None = None
         self.client: IecClient | None = None
 
+    @staticmethod
+    @callback
+    def async_get_options_flow(
+        config_entry: config_entries.ConfigEntry,
+    ) -> config_entries.OptionsFlow:
+        """Create the options flow."""
+        return IecOptionsFlow()
+
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
@@ -311,3 +323,147 @@ async def async_step_reauth_confirm(
             data_schema=vol.Schema(schema),
             errors=errors,
         )
+
+
+class IecOptionsFlow(config_entries.OptionsFlow):
+    """Handle IEC options for Gmail OTP configuration."""
+
+    def __init__(self) -> None:
+        """Initialize options flow."""
+        self._gmail_flow: Any = None  # GmailOAuthFlow instance
+        self._auth_url: str | None = None
+
+    async def async_step_init(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Manage the main options."""
+        errors: dict[str, str] = {}
+
+        if user_input is not None:
+            gmail_enabled = user_input.get(CONF_GMAIL_AUTO_REAUTH, False)
+
+            # If enabling Gmail and not already configured, go to credentials step
+            if gmail_enabled and not self.config_entry.options.get(CONF_GMAIL_TOKEN):
+                return await self.async_step_gmail_credentials()
+
+            # If disabling Gmail, clear the stored credentials
+            if not gmail_enabled:
+                return self.async_create_entry(
+                    title="",
+                    data={
+                        CONF_GMAIL_AUTO_REAUTH: False,
+                        CONF_GMAIL_TOKEN: None,
+                    },
+                )
+
+            # Gmail already configured, just save the enabled state
+            return self.async_create_entry(
+                title="",
+                data={
+                    **self.config_entry.options,
+                    CONF_GMAIL_AUTO_REAUTH: gmail_enabled,
+                },
+            )
+
+        # Check if Gmail is currently configured
+        gmail_configured = bool(self.config_entry.options.get(CONF_GMAIL_TOKEN))
+        gmail_enabled = self.config_entry.options.get(CONF_GMAIL_AUTO_REAUTH, False)
+
+        return self.async_show_form(
+            step_id="init",
+            data_schema=vol.Schema(
+                {
+                    vol.Optional(
+                        CONF_GMAIL_AUTO_REAUTH,
+                        default=gmail_enabled,
+                    ): bool,
+                }
+            ),
+            description_placeholders={
+                "gmail_status": "✅ Configured"
+                if gmail_configured
+                else "❌ Not configured"
+            },
+            errors=errors,
+        )
+
+    async def async_step_gmail_credentials(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle Gmail credentials input."""
+        errors: dict[str, str] = {}
+
+        if user_input is not None:
+            credentials_json = user_input.get(CONF_GMAIL_CREDENTIALS, "")
+
+            try:
+                # Validate JSON format
+                json.loads(credentials_json)
+
+                # Import here to avoid circular imports and optional dependency issues
+                from .gmail_otp_helper import GmailOAuthFlow
+
+                # Create OAuth flow and get authorization URL
+                self._gmail_flow = GmailOAuthFlow(credentials_json)
+                self._auth_url = await self.hass.async_add_executor_job(
+                    self._gmail_flow.get_authorization_url
+                )
+
+                return await self.async_step_gmail_authorize()
+
+            except json.JSONDecodeError:
+                errors["base"] = "invalid_credentials_json"
+            except Exception as err:  # noqa: BLE001
+                _LOGGER.exception("Error processing Gmail credentials: %s", err)
+                errors["base"] = "gmail_credentials_error"
+
+        return self.async_show_form(
+            step_id="gmail_credentials",
+            data_schema=vol.Schema(
+                {
+                    vol.Required(CONF_GMAIL_CREDENTIALS): str,
+                }
+            ),
+            errors=errors,
+        )
+
+    async def async_step_gmail_authorize(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
+        """Handle Gmail authorization code input."""
+        errors: dict[str, str] = {}
+
+        if user_input is not None:
+            auth_code = user_input.get("auth_code", "").strip()
+
+            try:
+                # Exchange code for tokens
+                token_data = await self.hass.async_add_executor_job(
+                    self._gmail_flow.exchange_code, auth_code
+                )
+
+                # Save tokens to options
+                return self.async_create_entry(
+                    title="",
+                    data={
+                        CONF_GMAIL_AUTO_REAUTH: True,
+                        CONF_GMAIL_TOKEN: token_data,
+                    },
+                )
+
+            except Exception as err:  # noqa: BLE001
+                _LOGGER.exception("Error exchanging Gmail auth code: %s", err)
+                errors["base"] = "gmail_auth_failed"
+
+        return self.async_show_form(
+            step_id="gmail_authorize",
+            data_schema=vol.Schema(
+                {
+                    vol.Required("auth_code"): str,
+                }
+            ),
+            description_placeholders={
+                "auth_url": self._auth_url or "",
+            },
+            errors=errors,
+        )

custom_components/iec/const.py

@@ -72,3 +72,8 @@
 ELECTRIC_INVOICE_DOC_ID = "1"
 ACCESS_TOKEN_ISSUED_AT = "iat"
 ACCESS_TOKEN_EXPIRATION_TIME = "exp"
+
+# Gmail OTP automation
+CONF_GMAIL_AUTO_REAUTH = "gmail_auto_reauth"
+CONF_GMAIL_CREDENTIALS = "gmail_credentials"
+CONF_GMAIL_TOKEN = "gmail_token"

custom_components/iec/coordinator.py

@@ -48,6 +48,8 @@
     ACCESS_TOKEN_ISSUED_AT,
     ATTRIBUTES_DICT_NAME,
     CONF_BP_NUMBER,
+    CONF_GMAIL_AUTO_REAUTH,
+    CONF_GMAIL_TOKEN,
     CONF_SELECTED_CONTRACTS,
     CONF_USER_ID,
     CONTRACT_DICT_NAME,
@@ -122,6 +124,7 @@ def __init__(
             session=self._api_session,
         )
         self._first_load: bool = True
+        self._gmail_auto_reauth_in_progress: bool = False
 
         @callback
         def _dummy_listener() -> None:
@@ -373,6 +376,84 @@ async def _get_power_size(self, connection_size) -> float:
                 )
         return power_size or 0.0
 
+    def _is_gmail_auto_reauth_enabled(self) -> bool:
+        """Check if Gmail auto-reauth is enabled and configured."""
+        options = self._config_entry.options
+        return bool(
+            options.get(CONF_GMAIL_AUTO_REAUTH) and options.get(CONF_GMAIL_TOKEN)
+        )
+
+    async def _attempt_gmail_auto_reauth(self) -> bool:
+        """Attempt automatic reauthentication using Gmail OTP.
+
+        Returns:
+            True if reauth succeeded, False otherwise.
+
+        """
+        if self._gmail_auto_reauth_in_progress:
+            _LOGGER.debug("Gmail auto-reauth already in progress, skipping")
+            return False
+
+        if not self._is_gmail_auto_reauth_enabled():
+            _LOGGER.debug("Gmail auto-reauth not enabled")
+            return False
+
+        self._gmail_auto_reauth_in_progress = True
+        _LOGGER.info("Attempting automatic reauthentication via Gmail OTP")
+
+        try:
+            from datetime import datetime as dt
+
+            from .gmail_otp_helper import GmailOtpHelper, GmailTokenRevoked
+
+            # Step 1: Request OTP from IEC
+            _LOGGER.debug("Requesting OTP from IEC API")
+            request_time = dt.now()
+            await self.api.login_with_id()
+
+            # Step 2: Poll Gmail for OTP
+            token_data = self._config_entry.options.get(CONF_GMAIL_TOKEN)
+            if not token_data:
+                _LOGGER.error("Gmail token data not found in options")
+                return False
+
+            gmail_helper = GmailOtpHelper(self.hass, token_data)
+
+            try:
+                otp = await gmail_helper.fetch_otp(request_time)
+            except GmailTokenRevoked:
+                _LOGGER.error("Gmail token revoked, manual reconfiguration required")
+                return False
+
+            if not otp:
+                _LOGGER.error("Failed to fetch OTP from Gmail within timeout")
+                return False
+
+            _LOGGER.debug("Successfully retrieved OTP from Gmail")
+
+            # Step 3: Verify OTP with IEC
+            await self.api.verify_otp(otp)
+
+            # Step 4: Update stored token
+            new_token = self.api.get_token()
+            new_data = {**self._entry_data, CONF_API_TOKEN: new_token.to_dict()}
+            self.hass.config_entries.async_update_entry(
+                entry=self._config_entry, data=new_data
+            )
+            self._entry_data = new_data
+
+            _LOGGER.info("Gmail auto-reauth completed successfully")
+            return True
+
+        except IECError as err:
+            _LOGGER.error("IEC API error during Gmail auto-reauth: %s", err)
+            return False
+        except Exception as err:  # noqa: BLE001
+            _LOGGER.exception("Unexpected error during Gmail auto-reauth: %s", err)
+            return False
+        finally:
+            self._gmail_auto_reauth_in_progress = False
+
     async def _get_readings(
         self,
         contract_id: int,
@@ -826,7 +907,23 @@ async def _async_update_data(
                     entry=self._config_entry, data=new_data
                 )
         except IECError as err:
-            raise ConfigEntryAuthFailed from err
+            # Attempt Gmail auto-reauth before failing
+            if self._is_gmail_auto_reauth_enabled():
+                _LOGGER.info("Token check failed, attempting Gmail auto-reauth")
+                if await self._attempt_gmail_auto_reauth():
+                    # Reauth succeeded, retry token check
+                    try:
+                        await self.api.check_token()
+                    except IECError as retry_err:
+                        _LOGGER.error("Token check failed after Gmail auto-reauth")
+                        raise ConfigEntryAuthFailed from retry_err
+                else:
+                    _LOGGER.error(
+                        "Gmail auto-reauth failed, falling back to manual reauth"
+                    )
+                    raise ConfigEntryAuthFailed from err
+            else:
+                raise ConfigEntryAuthFailed from err
 
         try:
             return await self._update_data()