supervisor/hassio/api/proxy.py

"""Utils for Home Assistant Proxy."""
import asyncio
from contextlib import asynccontextmanager
import logging

import aiohttp
from aiohttp import web
from aiohttp.web_exceptions import HTTPBadGateway, HTTPUnauthorized
from aiohttp.client_exceptions import ClientConnectorError
from aiohttp.hdrs import CONTENT_TYPE, AUTHORIZATION
import async_timeout

from ..const import HEADER_HA_ACCESS
from ..coresys import CoreSysAttributes
from ..exceptions import HomeAssistantAuthError, HomeAssistantAPIError, APIError

_LOGGER: logging.Logger = logging.getLogger(__name__)


class APIProxy(CoreSysAttributes):
    """API Proxy for Home Assistant."""

    def _check_access(self, request: web.Request):
        """Check the Hass.io token."""
        if AUTHORIZATION in request.headers:
            bearer = request.headers[AUTHORIZATION]
            hassio_token = bearer.split(" ")[-1]
        else:
            hassio_token = request.headers.get(HEADER_HA_ACCESS)

        addon = self.sys_addons.from_token(hassio_token)
        if not addon:
            _LOGGER.warning("Unknown Home Assistant API access!")
        elif not addon.access_homeassistant_api:
            _LOGGER.warning("Not permitted API access: %s", addon.slug)
        else:
            _LOGGER.debug("%s access from %s", request.path, addon.slug)
            return

        raise HTTPUnauthorized()

    @asynccontextmanager
    async def _api_client(self, request: web.Request, path: str, timeout: int = 300):
        """Return a client request with proxy origin for Home Assistant."""
        try:
            # read data
            with async_timeout.timeout(30):
                data = await request.read()

            if data:
                content_type = request.content_type
            else:
                content_type = None

            async with self.sys_homeassistant.make_request(
                request.method.lower(),
                f"api/{path}",
                content_type=content_type,
                data=data,
                timeout=timeout,
                params=request.query,
            ) as resp:
                yield resp
                return

        except HomeAssistantAuthError:
            _LOGGER.error("Authenticate error on API for request %s", path)
        except HomeAssistantAPIError:
            _LOGGER.error("Error on API for request %s", path)
        except aiohttp.ClientError as err:
            _LOGGER.error("Client error on API %s request %s", path, err)
        except asyncio.TimeoutError:
            _LOGGER.error("Client timeout error on API request %s", path)

        raise HTTPBadGateway()

    async def stream(self, request: web.Request):
        """Proxy HomeAssistant EventStream Requests."""
        self._check_access(request)

        _LOGGER.info("Home Assistant EventStream start")
        async with self._api_client(request, "stream", timeout=None) as client:
            response = web.StreamResponse()
            response.content_type = request.headers.get(CONTENT_TYPE)
            try:
                await response.prepare(request)
                async for data in client.content:
                    await response.write(data)

            except (aiohttp.ClientError, aiohttp.ClientPayloadError):
                pass

            _LOGGER.info("Home Assistant EventStream close")
            return response

    async def api(self, request: web.Request):
        """Proxy Home Assistant API Requests."""
        self._check_access(request)

        # Normal request
        path = request.match_info.get("path", "")
        async with self._api_client(request, path) as client:
            data = await client.read()
            return web.Response(
                body=data, status=client.status, content_type=client.content_type
            )

    async def _websocket_client(self):
        """Initialize a WebSocket API connection."""
        url = f"{self.sys_homeassistant.api_url}/api/websocket"

        try:
            client = await self.sys_websession_ssl.ws_connect(
                url, heartbeat=30, verify_ssl=False
            )

            # Handle authentication
            data = await client.receive_json()

            if data.get("type") == "auth_ok":
                return client

            if data.get("type") != "auth_required":
                # Invalid protocol
                _LOGGER.error("Got unexpected response from HA WebSocket: %s", data)
                raise APIError()

            if self.sys_homeassistant.refresh_token:
                await self.sys_homeassistant.ensure_access_token()
                await client.send_json(
                    {
                        "type": "auth",
                        "access_token": self.sys_homeassistant.access_token,
                    }
                )
            else:
                await client.send_json(
                    {
                        "type": "auth",
                        "api_password": self.sys_homeassistant.api_password,
                    }
                )

            data = await client.receive_json()

            if data.get("type") == "auth_ok":
                return client

            # Renew the Token is invalid
            if (
                data.get("type") == "invalid_auth"
                and self.sys_homeassistant.refresh_token
            ):
                self.sys_homeassistant.access_token = None
                return await self._websocket_client()

            raise HomeAssistantAuthError()

        except (RuntimeError, ValueError, ClientConnectorError) as err:
            _LOGGER.error("Client error on WebSocket API %s.", err)
        except HomeAssistantAuthError:
            _LOGGER.error("Failed authentication to Home Assistant WebSocket")

        raise APIError()

    async def websocket(self, request: web.Request):
        """Initialize a WebSocket API connection."""
        _LOGGER.info("Home Assistant WebSocket API request initialize")

        # init server
        server = web.WebSocketResponse(heartbeat=30)
        await server.prepare(request)

        # handle authentication
        try:
            await server.send_json(
                {"type": "auth_required", "ha_version": self.sys_homeassistant.version}
            )

            # Check API access
            response = await server.receive_json()
            hassio_token = response.get("api_password") or response.get("access_token")
            addon = self.sys_addons.from_token(hassio_token)

            if not addon or not addon.access_homeassistant_api:
                _LOGGER.warning("Unauthorized WebSocket access!")
                await server.send_json(
                    {"type": "auth_invalid", "message": "Invalid access"}
                )
                return server

            _LOGGER.info("WebSocket access from %s", addon.slug)

            await server.send_json(
                {"type": "auth_ok", "ha_version": self.sys_homeassistant.version}
            )
        except (RuntimeError, ValueError) as err:
            _LOGGER.error("Can't initialize handshake: %s", err)
            return server

        # init connection to hass
        try:
            client = await self._websocket_client()
        except APIError:
            return server

        _LOGGER.info("Home Assistant WebSocket API request running")
        try:
            client_read = None
            server_read = None
            while not server.closed and not client.closed:
                if not client_read:
                    client_read = self.sys_create_task(client.receive_str())
                if not server_read:
                    server_read = self.sys_create_task(server.receive_str())

                # wait until data need to be processed
                await asyncio.wait(
                    [client_read, server_read], return_when=asyncio.FIRST_COMPLETED
                )

                # server
                if server_read.done() and not client.closed:
                    server_read.exception()
                    await client.send_str(server_read.result())
                    server_read = None

                # client
                if client_read.done() and not server.closed:
                    client_read.exception()
                    await server.send_str(client_read.result())
                    client_read = None

        except asyncio.CancelledError:
            pass

        except (RuntimeError, ConnectionError, TypeError) as err:
            _LOGGER.info("Home Assistant WebSocket API error: %s", err)

        finally:
            if client_read:
                client_read.cancel()
            if server_read:
                server_read.cancel()

            # close connections
            if not client.closed:
                await client.close()
            if not server.closed:
                await server.close()

        _LOGGER.info("Home Assistant WebSocket API connection is closed")
        return server