pyh0n/pyhon/connection/handler.py

import json
from contextlib import asynccontextmanager

import aiohttp

from pyhon import const
from pyhon.connection.auth import HonAuth, _LOGGER
from pyhon.connection.device import HonDevice
from pyhon.exceptions import HonAuthenticationError


class HonBaseConnectionHandler:
    _HEADERS = {"user-agent": const.USER_AGENT, "Content-Type": "application/json"}

    def __init__(self, session=None):
        self._create_session = session is None
        self._session = session
        self._auth = None

    async def __aenter__(self):
        return await self.create()

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()

    async def create(self):
        if self._create_session:
            self._session = aiohttp.ClientSession()
        return self

    @asynccontextmanager
    async def _intercept(self, method, *args, loop=0, **kwargs):
        raise NotImplementedError

    @asynccontextmanager
    async def get(self, *args, **kwargs):
        async with self._intercept(self._session.get, *args, **kwargs) as response:
            yield response

    @asynccontextmanager
    async def post(self, *args, **kwargs):
        async with self._intercept(self._session.post, *args, **kwargs) as response:
            yield response

    async def close(self):
        if self._create_session:
            await self._session.close()


class HonConnectionHandler(HonBaseConnectionHandler):
    def __init__(self, email, password, session=None):
        super().__init__(session=session)
        self._device = HonDevice()
        self._email = email
        self._password = password
        if not self._email:
            raise HonAuthenticationError("An email address must be specified")
        if not self._password:
            raise HonAuthenticationError("A password address must be specified")

    @property
    def device(self):
        return self._device

    async def create(self):
        await super().create()
        self._auth = HonAuth(self._session, self._email, self._password, self._device)
        return self

    async def _check_headers(self, headers):
        if not (self._auth.cognito_token and self._auth.id_token):
            await self._auth.authenticate()
        headers["cognito-token"] = self._auth.cognito_token
        headers["id-token"] = self._auth.id_token
        return self._HEADERS | headers

    @asynccontextmanager
    async def _intercept(self, method, *args, loop=0, **kwargs):
        kwargs["headers"] = await self._check_headers(kwargs.get("headers", {}))
        async with method(*args, **kwargs) as response:
            if response.status in [401, 403] and loop == 0:
                _LOGGER.info("Try refreshing token...")
                await self._auth.refresh()
                async with self._intercept(
                    method, *args, loop=loop + 1, **kwargs
                ) as result:
                    yield result
            elif response.status in [401, 403] and loop == 1:
                _LOGGER.warning(
                    "%s - Error %s - %s",
                    response.request_info.url,
                    response.status,
                    await response.text(),
                )
                await self.create()
                async with self._intercept(
                    method, *args, loop=loop + 1, **kwargs
                ) as result:
                    yield result
            elif loop >= 2:
                _LOGGER.error(
                    "%s - Error %s - %s",
                    response.request_info.url,
                    response.status,
                    await response.text(),
                )
                raise HonAuthenticationError("Login failure")
            else:
                try:
                    await response.json()
                    yield response
                except json.JSONDecodeError:
                    _LOGGER.warning(
                        "%s - JsonDecodeError %s - %s",
                        response.request_info.url,
                        response.status,
                        await response.text(),
                    )
                    yield {}


class HonAnonymousConnectionHandler(HonBaseConnectionHandler):
    _HEADERS = HonBaseConnectionHandler._HEADERS | {"x-api-key": const.API_KEY}

    @asynccontextmanager
    async def _intercept(self, method, *args, loop=0, **kwargs):
        kwargs["headers"] = kwargs.pop("headers", {}) | self._HEADERS
        async with method(*args, **kwargs) as response:
            if response.status == 403:
                _LOGGER.error("Can't authenticate anymore")
            yield response