async def _handle_auth(self, ws_client: WebSocketResponse, token: str):
"""Handle authentication with JWT token."""
token_info = jwt.decode(token,
self.mass.web.jwt_key,
algorithms=["HS256"])
if self.mass.config.security.is_token_revoked(token_info):
raise AuthenticationError("Token is revoked")
ws_client.authenticated = True
self.mass.config.security.set_last_login(token_info["client_id"])
# TODO: store token/app_id on ws_client obj and periodiclaly check if token is expired or revoked
await self._send_json(ws_client, result="auth", data=token_info)
async def _websocket_handler(self, request: web.Request):
"""Handle websocket client."""
ws_client = WebSocketResponse()
ws_client.authenticated = False
await ws_client.prepare(request)
request.app["clients"].append(ws_client)
# handle incoming messages
async for msg in ws_client:
try:
if msg.type == WSMsgType.error:
LOGGER.warning("ws connection closed with exception %s",
ws_client.exception())
if msg.type != WSMsgType.text:
continue
if msg.data == "close":
await ws_client.close()
break
# regular message
json_msg = msg.json(loads=ujson.loads)
if "command" in json_msg and "data" in json_msg:
# handle command
await self._handle_command(
ws_client,
json_msg["command"],
json_msg["data"],
json_msg.get("id"),
)
elif "event" in json_msg:
# handle event
await self._handle_event(ws_client, json_msg["event"],
json_msg.get("data"))
except AuthenticationError as exc: # pylint:disable=broad-except
# disconnect client on auth errors
await self._send_json(ws_client, error=str(exc), **json_msg)
await ws_client.close(message=str(exc).encode())
except Exception as exc: # pylint:disable=broad-except
# log the error only
await self._send_json(ws_client, error=str(exc), **json_msg)
LOGGER.error("Error with WS client", exc_info=exc)
# websocket disconnected
request.app["clients"].remove(ws_client)
LOGGER.debug("websocket connection closed: %s", request.remote)
return ws_client
