|
7 | 7 |
|
8 | 8 | import aiohttp |
9 | 9 | from aiohttp import WSCloseCode, WSMessageTypeError, web |
| 10 | +from aiohttp.client_exceptions import ClientConnectorError |
10 | 11 | from aiohttp.client_ws import ClientWebSocketResponse |
11 | 12 | from aiohttp.hdrs import AUTHORIZATION, CONTENT_TYPE |
12 | 13 | from aiohttp.http_websocket import WSMsgType |
@@ -178,16 +179,57 @@ async def api(self, request: web.Request): |
178 | 179 |
|
179 | 180 | async def _websocket_client(self) -> ClientWebSocketResponse: |
180 | 181 | """Initialize a WebSocket API connection.""" |
| 182 | + url = f"{self.sys_homeassistant.api_url}/api/websocket" |
| 183 | + |
181 | 184 | try: |
182 | | - ws_client = await self.sys_homeassistant.api.connect_websocket( |
183 | | - max_msg_size=MAX_MESSAGE_SIZE_FROM_CORE |
| 185 | + client = await self.sys_websession.ws_connect( |
| 186 | + url, heartbeat=30, ssl=False, max_msg_size=MAX_MESSAGE_SIZE_FROM_CORE |
184 | 187 | ) |
185 | | - return ws_client.client |
186 | | - except HomeAssistantAPIError as err: |
187 | | - raise APIError( |
188 | | - f"Error connecting to Home Assistant WebSocket: {err}", |
189 | | - _LOGGER.error, |
190 | | - ) from err |
| 188 | + |
| 189 | + # Handle authentication |
| 190 | + data = await client.receive_json() |
| 191 | + |
| 192 | + if data.get("type") == "auth_ok": |
| 193 | + return client |
| 194 | + |
| 195 | + if data.get("type") != "auth_required": |
| 196 | + # Invalid protocol |
| 197 | + raise APIError( |
| 198 | + f"Got unexpected response from Home Assistant WebSocket: {data}", |
| 199 | + _LOGGER.error, |
| 200 | + ) |
| 201 | + |
| 202 | + # Auth session |
| 203 | + await self.sys_homeassistant.api.ensure_access_token() |
| 204 | + await client.send_json( |
| 205 | + { |
| 206 | + "type": "auth", |
| 207 | + "access_token": self.sys_homeassistant.api.access_token, |
| 208 | + }, |
| 209 | + dumps=json_dumps, |
| 210 | + ) |
| 211 | + |
| 212 | + data = await client.receive_json() |
| 213 | + |
| 214 | + if data.get("type") == "auth_ok": |
| 215 | + return client |
| 216 | + |
| 217 | + # Renew the Token is invalid |
| 218 | + if ( |
| 219 | + data.get("type") == "invalid_auth" |
| 220 | + and self.sys_homeassistant.refresh_token |
| 221 | + ): |
| 222 | + self.sys_homeassistant.api.access_token = None |
| 223 | + return await self._websocket_client() |
| 224 | + |
| 225 | + raise HomeAssistantAuthError() |
| 226 | + |
| 227 | + except (RuntimeError, ValueError, TypeError, ClientConnectorError) as err: |
| 228 | + _LOGGER.error("Client error on WebSocket API %s.", err) |
| 229 | + except HomeAssistantAuthError: |
| 230 | + _LOGGER.error("Failed authentication to Home Assistant WebSocket") |
| 231 | + |
| 232 | + raise APIError() |
191 | 233 |
|
192 | 234 | async def _proxy_message( |
193 | 235 | self, |
|
0 commit comments