Premier essai (NE PAS MERGER!!!)

pyhilo/api.py

@@ -27,6 +27,7 @@
     API_NOTIFICATIONS_ENDPOINT,
     API_REGISTRATION_ENDPOINT,
     API_REGISTRATION_HEADERS,
+    AUTOMATION_CHALLENGE_ENDPOINT,
     AUTOMATION_DEVICEHUB_ENDPOINT,
     DEFAULT_STATE_FILE,
     DEFAULT_USER_AGENT,
@@ -51,7 +52,7 @@
     get_state,
     set_state,
 )
-from pyhilo.websocket import WebsocketClient
+from pyhilo.websocket import WebsocketClient, WebsocketConfig, WebsocketManager
 
 
 class API:
@@ -216,17 +217,24 @@ async def _async_request(
         :rtype: dict[str, Any]
         """
         kwargs.setdefault("headers", self.headers)
+        access_token = await self.async_get_access_token()
+
         if endpoint.startswith(API_REGISTRATION_ENDPOINT):
             kwargs["headers"] = {**kwargs["headers"], **API_REGISTRATION_HEADERS}
         if endpoint.startswith(FB_INSTALL_ENDPOINT):
             kwargs["headers"] = {**kwargs["headers"], **FB_INSTALL_HEADERS}
         if endpoint.startswith(ANDROID_CLIENT_ENDPOINT):
             kwargs["headers"] = {**kwargs["headers"], **ANDROID_CLIENT_HEADERS}
         if host == API_HOSTNAME:
-            access_token = await self.async_get_access_token()
             kwargs["headers"]["authorization"] = f"Bearer {access_token}"
         kwargs["headers"]["Host"] = host
 
+        # ic-dev21 trying Leicas suggestion
+        if endpoint.startswith(AUTOMATION_CHALLENGE_ENDPOINT):
+            # remove Ocp-Apim-Subscription-Key header to avoid 401 error
+            kwargs["headers"].pop("Ocp-Apim-Subscription-Key", None)
+            kwargs["headers"]["authorization"] = f"Bearer {access_token}"
+
         data: dict[str, Any] = {}
         url = parse.urljoin(f"https://{host}", endpoint)
         if self.log_traces:
@@ -303,8 +311,9 @@ async def _async_handle_on_backoff(self, _: dict[str, Any]) -> None:
                 LOG.info(
                     "401 detected on websocket, refreshing websocket token. Old url: {self.ws_url} Old Token: {self.ws_token}"
                 )
+                LOG.info(f"401 detected on {err.request_info.url}")
                 async with self._backoff_refresh_lock_ws:
-                    (self.ws_url, self.ws_token) = await self.post_devicehub_negociate()
+                    await self.refresh_ws_token()
                     await self.get_websocket_params()
                 return
 
@@ -354,12 +363,23 @@ async def _async_post_init(self) -> None:
         LOG.debug("Websocket postinit")
         await self._get_fid()
         await self._get_device_token()
-        await self.refresh_ws_token()
-        self.websocket = WebsocketClient(self)
+        # await self.refresh_ws_token()
+        # self.websocket = WebsocketClient(self)
+
+        # Initialize WebsocketManager ic-dev21
+        self.websocket_manager = WebsocketManager(
+            self.session, self.async_request, self._state_yaml, set_state
+        )
+        await self.websocket_manager.initialize_websockets()
+
+        # Create both websocket clients
+        self.websocket = WebsocketClient(self.websocket_manager.devicehub)
+        self.websocket2 = WebsocketClient(self.websocket_manager.challengehub)
 
     async def refresh_ws_token(self) -> None:
-        (self.ws_url, self.ws_token) = await self.post_devicehub_negociate()
-        await self.get_websocket_params()
+        """Refresh the websocket token."""
+        await self.websocket_manager.refresh_token(self.websocket_manager.devicehub)
+        await self.websocket_manager.refresh_token(self.websocket_manager.challengehub)
 
     async def post_devicehub_negociate(self) -> tuple[str, str]:
         LOG.debug("Getting websocket url")

pyhilo/const.py

@@ -42,6 +42,8 @@
 
 # Automation server constant
 AUTOMATION_DEVICEHUB_ENDPOINT: Final = "/DeviceHub"
+AUTOMATION_CHALLENGE_ENDPOINT: Final = "/ChallengeHub"
+
 
 # Request constants
 DEFAULT_USER_AGENT: Final = f"PyHilo/{PYHILO_VERSION} HomeAssistant/{homeassistant.core.__version__} aiohttp/{aiohttp.__version__} Python/{platform.python_version()}"

pyhilo/websocket.py

@@ -7,17 +7,23 @@
 from enum import IntEnum
 import json
 from os import environ
-from typing import TYPE_CHECKING, Any, Callable, Dict
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple
+from urllib import parse
 
-from aiohttp import ClientWebSocketResponse, WSMsgType
+from aiohttp import ClientSession, ClientWebSocketResponse, WSMsgType
 from aiohttp.client_exceptions import (
     ClientError,
     ServerDisconnectedError,
     WSServerHandshakeError,
 )
 from yarl import URL
 
-from pyhilo.const import DEFAULT_USER_AGENT, LOG
+from pyhilo.const import (
+    AUTOMATION_CHALLENGE_ENDPOINT,
+    AUTOMATION_DEVICEHUB_ENDPOINT,
+    DEFAULT_USER_AGENT,
+    LOG,
+)
 from pyhilo.exceptions import (
     CannotConnectError,
     ConnectionClosedError,
@@ -376,3 +382,140 @@ async def async_invoke(
                 "type": inv_type,
             }
         )
+
+
+@dataclass
+class WebsocketConfig:
+    """Configuration for a websocket connection"""
+
+    endpoint: str
+    url: Optional[str] = None
+    token: Optional[str] = None
+    connection_id: Optional[str] = None
+    full_url: Optional[str] = None
+
+
+class WebsocketManager:
+    """Manages multiple websocket connections for the Hilo API"""
+
+    def __init__(
+        self, session: ClientSession, async_request, state_yaml: str, set_state_callback
+    ) -> None:
+        """Initialize the websocket manager.
+
+        Args:
+            session: The aiohttp client session
+            async_request: The async request method from the API class
+            state_yaml: Path to the state file
+            set_state_callback: Callback to save state
+        """
+        self.session = session
+        self.async_request = async_request
+        self._state_yaml = state_yaml
+        self._set_state = set_state_callback
+        self._shared_token = None  # ic-dev21 need to share the token
+
+        # Initialize websocket configurations
+        self.devicehub = WebsocketConfig(endpoint=AUTOMATION_DEVICEHUB_ENDPOINT)
+        self.challengehub = WebsocketConfig(endpoint=AUTOMATION_CHALLENGE_ENDPOINT)
+
+    async def initialize_websockets(self) -> None:
+        """Initialize both websocket connections"""
+        # ic-dev21 get token from device hub
+        await self.refresh_token(self.devicehub, get_new_token=True)
+        # ic-dev21 reuse it for challenge hub
+        await self.refresh_token(self.challengehub, get_new_token=True)
+
+    async def refresh_token(
+        self, config: WebsocketConfig, get_new_token: bool = True
+    ) -> None:
+        """Refresh token for a specific websocket configuration.
+
+        Args:
+            config: The websocket configuration to refresh
+        """
+        if get_new_token:
+            config.url, self._shared_token = await self._negotiate(config)
+            config.token = self._shared_token
+        else:
+            # ic-dev21 reuse existing token but get new URL
+            config.url, _ = await self._negotiate(config)
+            config.token = self._shared_token
+
+        await self._get_websocket_params(config)
+
+    async def _negotiate(self, config: WebsocketConfig) -> Tuple[str, str]:
+        """Negotiate websocket connection and get URL and token.
+
+        Args:
+            config: The websocket configuration to negotiate
+
+        Returns:
+            Tuple containing the websocket URL and access token
+        """
+        LOG.debug(f"Getting websocket url for {config.endpoint}")
+        url = f"{config.endpoint}/negotiate"
+        LOG.debug(f"Negotiate URL is {url}")
+
+        resp = await self.async_request("post", url)
+        ws_url = resp.get("url")
+        ws_token =   resp.get("accessToken")
+
+
+        # Save state
+        state_key = (
+            "websocket"
+            if config.endpoint == "AUTOMATION_DEVICEHUB_ENDPOINT"
+            else "websocket2"
+        )
+        await self._set_state(
+            self._state_yaml,
+            state_key,
+            {
+                "url": ws_url,
+                "token": ws_token,
+            },
+        )
+
+        return ws_url, ws_token
+
+    async def _get_websocket_params(self, config: WebsocketConfig) -> None:
+        """Get websocket parameters including connection ID.
+
+        Args:
+            config: The websocket configuration to get parameters for
+        """
+        uri = parse.urlparse(config.url)
+        LOG.debug(f"Getting websocket params for {config.endpoint}")
+        LOG.debug(f"Getting uri {uri}")
+
+        resp = await self.async_request(
+            "post",
+            f"{uri.path}negotiate?{uri.query}",
+            host=uri.netloc,
+            headers={
+                "authorization": f"Bearer {config.token}",
+            },
+        )
+
+        config.connection_id = resp.get("connectionId", "")
+        config.full_url = (
+            f"{config.url}&id={config.connection_id}&access_token={config.token}"
+        )
+        LOG.debug(f"Getting full ws URL {config.full_url}")
+
+        transport_dict = resp.get("availableTransports", [])
+        websocket_dict = {
+            "connection_id": config.connection_id,
+            "available_transports": transport_dict,
+            "full_ws_url": config.full_url,
+        }
+
+        # Save state
+        state_key = (
+            "websocket"
+            if config.endpoint == "AUTOMATION_DEVICEHUB_ENDPOINT"
+            else "websocket2"
+        )
+        LOG.debug(f"Calling set_state {state_key}_params")
+        await self._set_state(self._state_yaml, state_key, websocket_dict)