core/homeassistant/helpers/aiohttp_client.py

"""Helper for aiohttp webclient stuff."""
import asyncio
import sys
from ssl import SSLContext  # noqa: F401
from typing import Any, Awaitable, Optional, cast
from typing import Union  # noqa: F401

import aiohttp
from aiohttp.hdrs import USER_AGENT, CONTENT_TYPE
from aiohttp import web
from aiohttp.web_exceptions import HTTPGatewayTimeout, HTTPBadGateway
import async_timeout

from homeassistant.core import callback, Event
from homeassistant.const import EVENT_HOMEASSISTANT_CLOSE, __version__
from homeassistant.helpers.typing import HomeAssistantType
from homeassistant.loader import bind_hass
from homeassistant.util import ssl as ssl_util

DATA_CONNECTOR = 'aiohttp_connector'
DATA_CONNECTOR_NOTVERIFY = 'aiohttp_connector_notverify'
DATA_CLIENTSESSION = 'aiohttp_clientsession'
DATA_CLIENTSESSION_NOTVERIFY = 'aiohttp_clientsession_notverify'
SERVER_SOFTWARE = 'HomeAssistant/{0} aiohttp/{1} Python/{2[0]}.{2[1]}'.format(
    __version__, aiohttp.__version__, sys.version_info)


@callback
@bind_hass
def async_get_clientsession(hass: HomeAssistantType,
                            verify_ssl: bool = True) -> aiohttp.ClientSession:
    """Return default aiohttp ClientSession.

    This method must be run in the event loop.
    """
    if verify_ssl:
        key = DATA_CLIENTSESSION
    else:
        key = DATA_CLIENTSESSION_NOTVERIFY

    if key not in hass.data:
        hass.data[key] = async_create_clientsession(hass, verify_ssl)

    return cast(aiohttp.ClientSession, hass.data[key])


@callback
@bind_hass
def async_create_clientsession(hass: HomeAssistantType,
                               verify_ssl: bool = True,
                               auto_cleanup: bool = True,
                               **kwargs: Any) -> aiohttp.ClientSession:
    """Create a new ClientSession with kwargs, i.e. for cookies.

    If auto_cleanup is False, you need to call detach() after the session
    returned is no longer used. Default is True, the session will be
    automatically detached on homeassistant_stop.

    This method must be run in the event loop.
    """
    connector = _async_get_connector(hass, verify_ssl)

    clientsession = aiohttp.ClientSession(
        loop=hass.loop,
        connector=connector,
        headers={USER_AGENT: SERVER_SOFTWARE},
        **kwargs
    )

    if auto_cleanup:
        _async_register_clientsession_shutdown(hass, clientsession)

    return clientsession


@bind_hass
async def async_aiohttp_proxy_web(
        hass: HomeAssistantType, request: web.BaseRequest,
        web_coro: Awaitable[aiohttp.ClientResponse], buffer_size: int = 102400,
        timeout: int = 10) -> Optional[web.StreamResponse]:
    """Stream websession request to aiohttp web response."""
    try:
        with async_timeout.timeout(timeout, loop=hass.loop):
            req = await web_coro

    except asyncio.CancelledError:
        # The user cancelled the request
        return None

    except asyncio.TimeoutError as err:
        # Timeout trying to start the web request
        raise HTTPGatewayTimeout() from err

    except aiohttp.ClientError as err:
        # Something went wrong with the connection
        raise HTTPBadGateway() from err

    try:
        return await async_aiohttp_proxy_stream(
            hass,
            request,
            req.content,
            req.headers.get(CONTENT_TYPE)
        )
    finally:
        req.close()


@bind_hass
async def async_aiohttp_proxy_stream(hass: HomeAssistantType,
                                     request: web.BaseRequest,
                                     stream: aiohttp.StreamReader,
                                     content_type: str,
                                     buffer_size: int = 102400,
                                     timeout: int = 10) -> web.StreamResponse:
    """Stream a stream to aiohttp web response."""
    response = web.StreamResponse()
    response.content_type = content_type
    await response.prepare(request)

    try:
        while True:
            with async_timeout.timeout(timeout, loop=hass.loop):
                data = await stream.read(buffer_size)

            if not data:
                break
            await response.write(data)

    except (asyncio.TimeoutError, aiohttp.ClientError):
        # Something went wrong fetching data, closed connection
        pass

    return response


@callback
def _async_register_clientsession_shutdown(
        hass: HomeAssistantType, clientsession: aiohttp.ClientSession) -> None:
    """Register ClientSession close on Home Assistant shutdown.

    This method must be run in the event loop.
    """
    @callback
    def _async_close_websession(event: Event) -> None:
        """Close websession."""
        clientsession.detach()

    hass.bus.async_listen_once(
        EVENT_HOMEASSISTANT_CLOSE, _async_close_websession)


@callback
def _async_get_connector(hass: HomeAssistantType,
                         verify_ssl: bool = True) -> aiohttp.BaseConnector:
    """Return the connector pool for aiohttp.

    This method must be run in the event loop.
    """
    key = DATA_CONNECTOR if verify_ssl else DATA_CONNECTOR_NOTVERIFY

    if key in hass.data:
        return cast(aiohttp.BaseConnector, hass.data[key])

    if verify_ssl:
        ssl_context = \
            ssl_util.client_context()  # type: Union[bool, SSLContext]
    else:
        ssl_context = False

    connector = aiohttp.TCPConnector(loop=hass.loop,
                                     enable_cleanup_closed=True,
                                     ssl=ssl_context,
                                     )
    hass.data[key] = connector

    async def _async_close_connector(event: Event) -> None:
        """Close connector pool."""
        await connector.close()

    hass.bus.async_listen_once(
        EVENT_HOMEASSISTANT_CLOSE, _async_close_connector)

    return connector