Coverage for src/dev_shell/client.py: 0%

1"""Minimal Home Assistant websocket API client.""" 

2 

3from __future__ import annotations 

4 

5import itertools 

6import json 

7import os 

8from typing import Any, Self 

9from urllib.parse import urlsplit, urlunsplit 

10 

11import websockets 

12 

13 

14class DevShellError(Exception): 

15 """Connection, auth or command failure (not an error in the user's code).""" 

16 

17 

18def resolve_url(url: str | None) -> str: 

19 """Accept http(s)://host:8123 or ws(s)://... and return the websocket endpoint.""" 

20 if not url: 

21 # Inside a Home Assistant add-on (e.g. Studio Code Server) talk via the supervisor. 

22 if os.environ.get("SUPERVISOR_TOKEN"): 

23 return "ws://supervisor/core/websocket" 

24 url = "http://homeassistant.local:8123" 

25 parts = urlsplit(url) 

26 scheme = {"http": "ws", "https": "wss"}.get(parts.scheme, parts.scheme) 

27 path = parts.path.rstrip("/") 

28 if not path.endswith("/websocket"): 

29 path += "/api/websocket" 

30 return urlunsplit((scheme, parts.netloc, path, "", "")) 

31 

32 

33def resolve_token(token: str | None) -> str: 

34 token = token or os.environ.get("HASS_TOKEN") or os.environ.get("SUPERVISOR_TOKEN") 

35 if not token: 

36 raise DevShellError( 

37 "No access token: set HASS_TOKEN or pass --token " 

38 "(create one under Profile > Security > Long-lived access tokens)" 

39 ) 

40 return token 

41 

42 

43class Client: 

44 def __init__(self, url: str, token: str) -> None: 

45 self.url = url 

46 self._token = token 

47 self._ids = itertools.count(1) 

48 self._ws: Any = None 

49 

50 async def __aenter__(self) -> Self: 

51 try: 

52 self._ws = await websockets.connect(self.url, max_size=None) 

53 except (OSError, websockets.InvalidURI, websockets.InvalidHandshake) as err: 

54 raise DevShellError(f"Cannot connect to {self.url}: {err}") from err 

55 msg = await self._recv() 

56 if msg.get("type") != "auth_required": 

57 raise DevShellError(f"Unexpected greeting: {msg}") 

58 await self._ws.send(json.dumps({"type": "auth", "access_token": self._token})) 

59 msg = await self._recv() 

60 if msg.get("type") != "auth_ok": 

61 raise DevShellError(f"Authentication failed: {msg.get('message', msg)}") 

62 return self 

63 

64 async def __aexit__(self, *exc: object) -> None: 

65 await self._ws.close() 

66 

67 async def _recv(self) -> dict[str, Any]: 

68 try: 

69 return json.loads(await self._ws.recv()) 

70 except websockets.ConnectionClosed as err: 

71 raise DevShellError(f"Connection closed: {err}") from err 

72 

73 async def call(self, type_: str, **payload: Any) -> Any: 

74 msg_id = next(self._ids) 

75 await self._ws.send(json.dumps({"id": msg_id, "type": type_, **payload})) 

76 while True: 

77 msg = await self._recv() 

78 if msg.get("id") != msg_id or msg.get("type") != "result": 

79 continue 

80 if not msg["success"]: 

81 error = msg.get("error", {}) 

82 if error.get("code") == "unknown_command": 

83 raise DevShellError( 

84 f"{type_} not available: is the dev_shell_server integration " 

85 "installed and `dev_shell_server:` in configuration.yaml?" 

86 ) 

87 raise DevShellError(f"{type_} failed: {error.get('message', error)}") 

88 return msg["result"]