"""Module to handle messages from Home Assistant cloud.""" from __future__ import annotations import asyncio from contextlib import suppress import logging import pprint import random from typing import TYPE_CHECKING, Any import uuid from . import iot_base from .events.types import RelayerConnectedEvent, RelayerDisconnectedEvent from .utils import Registry if TYPE_CHECKING: from . import Cloud, _ClientT HANDLERS = Registry() _LOGGER = logging.getLogger(__name__) class UnknownHandler(Exception): """Exception raised when trying to handle unknown handler.""" class ErrorMessage(Exception): """Exception raised when there was error handling message in the cloud.""" def __init__(self, error: Any) -> None: """Initialize Error Message.""" super().__init__("Error in Cloud") self.error = error class HandlerError(Exception): """Exception raised when the handler failed.""" def __init__(self, error: str) -> None: """Initialize Error Message.""" super().__init__("Error in handler") self.error = error class CloudIoT(iot_base.BaseIoT): """Class to manage the IoT connection.""" mark_connected_after_first_message: bool = True def __init__(self, cloud: Cloud[_ClientT]) -> None: """Initialize the CloudIoT class.""" super().__init__(cloud) # Local code waiting for a response self._response_handler: dict[str, asyncio.Future[Any]] = {} # Register start/stop cloud.register_on_start(self.start) cloud.register_on_stop(self.disconnect) @property def package_name(self) -> str: """Return the package name for logging.""" return __name__ @property def ws_heartbeat(self) -> float | None: """Server to connect to.""" return 300 @property def ws_server_url(self) -> str: """Server to connect to.""" return self.cloud.service_discovery.action_url("relayer_connect") async def start(self) -> None: """Start the CloudIoT server.""" if self.cloud.subscription_expired: return asyncio.create_task(self.connect()) async def async_send_message( self, handler: str, payload: Any, expect_answer: bool = True, ) -> Any | None: """Send a message.""" msgid = uuid.uuid4().hex fut: asyncio.Future[Any] | None = None if expect_answer: fut = self._response_handler[msgid] = asyncio.Future() try: await self.async_send_json_message( {"msgid": msgid, "handler": handler, "payload": payload}, ) if expect_answer and fut is not None: return await fut return None finally: self._response_handler.pop(msgid, None) def async_handle_message(self, msg: dict[str, Any]) -> None: """Handle a message.""" response_handler = self._response_handler.get(msg["msgid"]) if response_handler is not None: if "payload" in msg: response_handler.set_result(msg["payload"]) else: response_handler.set_exception(ErrorMessage(msg["error"])) return asyncio.create_task(self._async_handle_handler_message(msg)) async def _async_handle_handler_message(self, message: dict[str, Any]) -> None: """Handle incoming IoT message.""" response = {"msgid": message["msgid"]} try: handler = HANDLERS.get(message["handler"]) if handler is None: raise UnknownHandler result = await handler(self.cloud, message.get("payload")) # No response from handler if result is None: return response["payload"] = result except UnknownHandler: response["error"] = "unknown-handler" except HandlerError as err: self._logger.warning("Error handling message: %s", err.error) response["error"] = err.error except Exception: # pylint: disable=broad-except self._logger.exception("Unexpected error handling message") response["error"] = "exception" # Client is unset in case the connection has been lost. if self.client is None: return if self._logger.isEnabledFor(logging.DEBUG): self._logger.debug("Publishing message:\n%s\n", pprint.pformat(response)) # Suppress when client is closing. with suppress(ConnectionResetError): await self.client.send_json(response) async def _connected(self) -> None: """Handle connected.""" await super()._connected() await self.cloud.events.publish(event=RelayerConnectedEvent()) await self.cloud.client.cloud_connected() async def _disconnected(self) -> None: """Handle connected.""" await super()._disconnected() await self.cloud.events.publish( event=RelayerDisconnectedEvent(reason=self.last_disconnect_reason) ) await self.cloud.client.cloud_disconnected() @HANDLERS.register("system") async def async_handle_system(cloud: Cloud[_ClientT], payload: dict[str, Any]) -> None: """Handle an incoming IoT message for System.""" return await cloud.client.async_system_message(payload) @HANDLERS.register("alexa") async def async_handle_alexa( cloud: Cloud[_ClientT], payload: dict[str, Any], ) -> dict[str, Any]: """Handle an incoming IoT message for Alexa.""" return await cloud.client.async_alexa_message(payload) @HANDLERS.register("google_actions") async def async_handle_google_actions( cloud: Cloud[_ClientT], payload: dict[str, Any], ) -> dict[str, Any]: """Handle an incoming IoT message for Google Actions.""" return await cloud.client.async_google_message(payload) @HANDLERS.register("cloud") async def async_handle_cloud(cloud: Cloud[_ClientT], payload: dict[str, Any]) -> None: """Handle an incoming IoT message for cloud component.""" action = payload["action"] if action == "logout": # Log out of Home Assistant Cloud await cloud.logout() _LOGGER.error( "You have been logged out from Home Assistant cloud: %s", payload["reason"], ) elif action == "disconnect_remote": # Disconnect Remote connection await cloud.remote.disconnect(clear_snitun_token=True) elif action == "evaluate_remote_security": async def _reconnect() -> None: """Reconnect after a random timeout.""" await asyncio.sleep(random.randint(60, 7200)) await cloud.remote.disconnect(clear_snitun_token=True) await cloud.remote.connect() # Reconnect to remote frontends cloud.client.loop.create_task(_reconnect()) elif action in ("user_notification", "critical_user_notification"): # Send user Notification cloud.client.user_message( "homeassistant_cloud_notification", payload["title"], payload["message"], ) else: _LOGGER.warning("Received unknown cloud action: %s", action) @HANDLERS.register("remote_sni") async def async_handle_remote_sni( cloud: Cloud[_ClientT], payload: dict[str, Any], # noqa: ARG001 ) -> dict[str, Any]: """Handle remote UI requests for cloud.""" await cloud.client.async_cloud_connect_update(True) return {"server": cloud.remote.snitun_server} @HANDLERS.register("connection_info") async def async_handle_connection_info( cloud: Cloud[_ClientT], payload: dict[str, Any], ) -> dict[str, Any]: """Handle connection info requests for cloud.""" return await cloud.client.async_cloud_connection_info(payload) @HANDLERS.register("webhook") async def async_handle_webhook( cloud: Cloud[_ClientT], payload: dict[str, Any], ) -> dict[str, Any]: """Handle an incoming IoT message for cloud webhooks.""" return await cloud.client.async_webhook_message(payload)