From b68351b693687033ab6c58a00ce6599c7c4e4e70 Mon Sep 17 00:00:00 2001 From: Elad Bar Date: Tue, 24 May 2022 16:33:13 +0300 Subject: [PATCH] Separation between the camera component and data provider --- .../xiaomi_cloud_map_extractor/camera.py | 339 ++---------------- .../xiaomi_cloud_map_extractor/const.py | 20 ++ .../xiaomi_cloud_map_extractor/enums.py | 18 + .../vacuum_manager.py | 338 +++++++++++++++++ 4 files changed, 411 insertions(+), 304 deletions(-) create mode 100644 custom_components/xiaomi_cloud_map_extractor/enums.py create mode 100644 custom_components/xiaomi_cloud_map_extractor/vacuum_manager.py diff --git a/custom_components/xiaomi_cloud_map_extractor/camera.py b/custom_components/xiaomi_cloud_map_extractor/camera.py index 0568631..75e8cb6 100644 --- a/custom_components/xiaomi_cloud_map_extractor/camera.py +++ b/custom_components/xiaomi_cloud_map_extractor/camera.py @@ -1,56 +1,27 @@ -import io import logging -import time +import voluptuous as vol + + +from abc import ABC from datetime import timedelta -from enum import Enum -from typing import Any, Dict, List, Optional +from typing import Optional -from custom_components.xiaomi_cloud_map_extractor.common.map_data import MapData -from custom_components.xiaomi_cloud_map_extractor.common.vacuum import XiaomiCloudVacuum -from custom_components.xiaomi_cloud_map_extractor.types import Colors, Drawables, ImageConfig, Sizes, Texts +from homeassistant.const import CONF_HOST, CONF_TOKEN, CONF_USERNAME, CONF_PASSWORD, CONF_NAME +from homeassistant.helpers import config_validation as cv +from homeassistant.helpers.event import async_track_time_interval + +from custom_components.xiaomi_cloud_map_extractor.vacuum_manager import VacuumManager -try: - from miio import RoborockVacuum, DeviceException -except ImportError: - from miio import Vacuum as RoborockVacuum, DeviceException -import PIL.Image as Image -import voluptuous as vol from homeassistant.components.camera import Camera, ENTITY_ID_FORMAT, PLATFORM_SCHEMA, SUPPORT_ON_OFF -from homeassistant.const import CONF_HOST, CONF_NAME, CONF_PASSWORD, CONF_TOKEN, CONF_USERNAME -from homeassistant.helpers import config_validation as cv from homeassistant.helpers.entity import generate_entity_id from homeassistant.helpers.reload import async_setup_reload_service -from custom_components.xiaomi_cloud_map_extractor.common.map_data_parser import MapDataParser -from custom_components.xiaomi_cloud_map_extractor.common.xiaomi_cloud_connector import XiaomiCloudConnector from custom_components.xiaomi_cloud_map_extractor.const import * -from custom_components.xiaomi_cloud_map_extractor.dreame.vacuum import DreameVacuum -from custom_components.xiaomi_cloud_map_extractor.roidmi.vacuum import RoidmiVacuum -from custom_components.xiaomi_cloud_map_extractor.unsupported.vacuum import UnsupportedVacuum -from custom_components.xiaomi_cloud_map_extractor.viomi.vacuum import ViomiVacuum -from custom_components.xiaomi_cloud_map_extractor.xiaomi.vacuum import XiaomiVacuum _LOGGER = logging.getLogger(__name__) SCAN_INTERVAL = timedelta(seconds=5) -DEFAULT_TRIMS = { - CONF_LEFT: 0, - CONF_RIGHT: 0, - CONF_TOP: 0, - CONF_BOTTOM: 0 -} - -DEFAULT_SIZES = { - CONF_SIZE_VACUUM_RADIUS: 6, - CONF_SIZE_PATH_WIDTH: 1, - CONF_SIZE_IGNORED_OBSTACLE_RADIUS: 3, - CONF_SIZE_IGNORED_OBSTACLE_WITH_PHOTO_RADIUS: 3, - CONF_SIZE_OBSTACLE_RADIUS: 3, - CONF_SIZE_OBSTACLE_WITH_PHOTO_RADIUS: 3, - CONF_SIZE_CHARGER_RADIUS: 6 -} - COLOR_SCHEMA = vol.Or( vol.All(vol.Length(min=3, max=3), vol.ExactSequence((cv.byte, cv.byte, cv.byte)), vol.Coerce(tuple)), vol.All(vol.Length(min=4, max=4), vol.ExactSequence((cv.byte, cv.byte, cv.byte, cv.byte)), vol.Coerce(tuple)) @@ -123,293 +94,53 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=None): await async_setup_reload_service(hass, DOMAIN, PLATFORMS) - host = config[CONF_HOST] - token = config[CONF_TOKEN] - username = config[CONF_USERNAME] - password = config[CONF_PASSWORD] - country = config[CONF_COUNTRY] - name = config[CONF_NAME] - should_poll = config[CONF_AUTO_UPDATE] - image_config = config[CONF_MAP_TRANSFORM] - colors = config[CONF_COLORS] - room_colors = config[CONF_ROOM_COLORS] - for room, color in room_colors.items(): - colors[f"{COLOR_ROOM_PREFIX}{room}"] = color - drawables = config[CONF_DRAW] - sizes = config[CONF_SIZES] - texts = config[CONF_TEXTS] - if DRAWABLE_ALL in drawables: - drawables = CONF_AVAILABLE_DRAWABLES[1:] - attributes = config[CONF_ATTRIBUTES] - store_map_raw = config[CONF_STORE_MAP_RAW] - store_map_image = config[CONF_STORE_MAP_IMAGE] - store_map_path = config[CONF_STORE_MAP_PATH] - force_api = config[CONF_FORCE_API] - entity_id = generate_entity_id(ENTITY_ID_FORMAT, name, hass=hass) - async_add_entities([VacuumCamera(entity_id, host, token, username, password, country, name, should_poll, - image_config, colors, drawables, sizes, texts, attributes, store_map_raw, - store_map_image, store_map_path, force_api)]) + manager = VacuumManager(config) + async_add_entities([VacuumCamera(hass, manager)]) -class VacuumCamera(Camera): - def __init__(self, entity_id: str, host: str, token: str, username: str, password: str, country: str, name: str, - should_poll: bool, image_config: ImageConfig, colors: Colors, drawables: Drawables, sizes: Sizes, - texts: Texts, attributes: List[str], store_map_raw: bool, store_map_image: bool, store_map_path: str, - force_api: str): + +class VacuumCamera(Camera, ABC): + def __init__(self, hass, manager: VacuumManager): super().__init__() - self.entity_id = entity_id + + self._hass = hass + self._manager = manager + self.entity_id = generate_entity_id(ENTITY_ID_FORMAT, manager.name, hass=hass) + self.content_type = CONTENT_TYPE - self._vacuum = RoborockVacuum(host, token) - self._connector = XiaomiCloudConnector(username, password) - self._status = CameraStatus.INITIALIZING - self._device = None - self._name = name - self._should_poll = should_poll - self._image_config = image_config - self._colors = colors - self._drawables = drawables - self._sizes = sizes - self._texts = texts - self._attributes = attributes - self._store_map_raw = store_map_raw - self._store_map_image = store_map_image - self._store_map_path = store_map_path - self._forced_api = force_api - self._used_api = None - self._map_saved = None - self._image = None - self._map_data = None - self._logged_in = False - self._logged_in_previously = True - self._received_map_name_previously = True - self._country = country + + self._remove_async_track_time = async_track_time_interval( + self._hass, self._manager.update, SCAN_INTERVAL + ) async def async_added_to_hass(self) -> None: self.async_schedule_update_ha_state(True) @property - def frame_interval(self) -> float: + def frame_interval(self): return 1 def camera_image(self, width: Optional[int] = None, height: Optional[int] = None) -> Optional[bytes]: - return self._image + return self._manager.image @property - def name(self) -> str: - return self._name + def name(self): + return self._manager.name def turn_on(self): - self._should_poll = True + self._manager.turn_on() def turn_off(self): - self._should_poll = False + self._manager.turn_off() @property - def supported_features(self) -> int: + def supported_features(self): return SUPPORT_ON_OFF @property - def extra_state_attributes(self) -> Dict[str, Any]: - attributes = {} - if self._map_data is not None: - attributes.update(self.extract_attributes(self._map_data, self._attributes, self._country)) - if self._store_map_raw: - attributes[ATTRIBUTE_MAP_SAVED] = self._map_saved - if self._device is not None: - attributes[ATTR_MODEL] = self._device.model - attributes[ATTR_USED_API] = self._used_api - if self._connector.two_factor_auth_url is not None: - attributes[ATTR_TWO_FACTOR_AUTH] = self._connector.two_factor_auth_url - return attributes + def extra_state_attributes(self): + return self._manager.attributes @property - def should_poll(self) -> bool: - return self._should_poll - - @staticmethod - def extract_attributes(map_data: MapData, attributes_to_return: List[str], country) -> Dict[str, Any]: - attributes = {} - rooms = [] - if map_data.rooms is not None: - rooms = dict(filter(lambda x: x[0] is not None, ((x[0], x[1].name) for x in map_data.rooms.items()))) - if len(rooms) == 0: - rooms = list(map_data.rooms.keys()) - for name, value in { - ATTRIBUTE_CALIBRATION: map_data.calibration(), - ATTRIBUTE_CHARGER: map_data.charger, - ATTRIBUTE_CLEANED_ROOMS: map_data.cleaned_rooms, - ATTRIBUTE_COUNTRY: country, - ATTRIBUTE_GOTO: map_data.goto, - ATTRIBUTE_GOTO_PATH: map_data.goto_path, - ATTRIBUTE_GOTO_PREDICTED_PATH: map_data.predicted_path, - ATTRIBUTE_IGNORED_OBSTACLES: map_data.ignored_obstacles, - ATTRIBUTE_IGNORED_OBSTACLES_WITH_PHOTO: map_data.ignored_obstacles_with_photo, - ATTRIBUTE_IMAGE: map_data.image, - ATTRIBUTE_IS_EMPTY: map_data.image.is_empty, - ATTRIBUTE_MAP_NAME: map_data.map_name, - ATTRIBUTE_NO_GO_AREAS: map_data.no_go_areas, - ATTRIBUTE_NO_MOPPING_AREAS: map_data.no_mopping_areas, - ATTRIBUTE_OBSTACLES: map_data.obstacles, - ATTRIBUTE_OBSTACLES_WITH_PHOTO: map_data.obstacles_with_photo, - ATTRIBUTE_PATH: map_data.path, - ATTRIBUTE_ROOM_NUMBERS: rooms, - ATTRIBUTE_ROOMS: map_data.rooms, - ATTRIBUTE_VACUUM_POSITION: map_data.vacuum_position, - ATTRIBUTE_VACUUM_ROOM: map_data.vacuum_room, - ATTRIBUTE_VACUUM_ROOM_NAME: map_data.vacuum_room_name, - ATTRIBUTE_WALLS: map_data.walls, - ATTRIBUTE_ZONES: map_data.zones - }.items(): - if name in attributes_to_return: - attributes[name] = value - return attributes - - def update(self): - counter = 10 - if self._status != CameraStatus.TWO_FACTOR_AUTH_REQUIRED and not self._logged_in: - self._handle_login() - if self._device is None and self._logged_in: - self._handle_device() - map_name = self._handle_map_name(counter) - if map_name == "retry" and self._device is not None: - self._status = CameraStatus.FAILED_TO_RETRIEVE_MAP_FROM_VACUUM - self._received_map_name_previously = map_name != "retry" - if self._logged_in and map_name != "retry" and self._device is not None: - self._handle_map_data(map_name) - else: - _LOGGER.debug("Unable to retrieve map, reasons: Logged in - %s, map name - %s, device retrieved - %s", - self._logged_in, map_name, self._device is not None) - self._set_map_data(MapDataParser.create_empty(self._colors, str(self._status))) - self._logged_in_previously = self._logged_in - - def _handle_login(self): - _LOGGER.debug("Logging in...") - self._logged_in = self._connector.login() - if self._logged_in is None: - _LOGGER.debug("2FA required") - self._status = CameraStatus.TWO_FACTOR_AUTH_REQUIRED - elif self._logged_in: - _LOGGER.debug("Logged in") - self._status = CameraStatus.LOGGED_IN - else: - _LOGGER.debug("Failed to log in") - self._status = CameraStatus.FAILED_LOGIN - if self._logged_in_previously: - _LOGGER.error("Unable to log in, check credentials") - - def _handle_device(self): - _LOGGER.debug("Retrieving device info, country: %s", self._country) - country, user_id, device_id, model = self._connector.get_device_details(self._vacuum.token, self._country) - if model is not None: - self._country = country - _LOGGER.debug("Retrieved device model: %s", model) - self._device = self._create_device(user_id, device_id, model) - _LOGGER.debug("Created device, used api: %s", self._used_api) - else: - _LOGGER.error("Failed to retrieve model") - self._status = CameraStatus.FAILED_TO_RETRIEVE_DEVICE - - def _handle_map_name(self, counter: int) -> str: - map_name = "retry" - if self._device is not None and not self._device.should_get_map_from_vacuum(): - map_name = "0" - while map_name == "retry" and counter > 0: - _LOGGER.debug("Retrieving map name from device") - time.sleep(0.1) - try: - map_name = self._vacuum.map()[0] - _LOGGER.debug("Map name %s", map_name) - except OSError as exc: - _LOGGER.error("Got OSError while fetching the state: %s", exc) - except DeviceException as exc: - if self._received_map_name_previously: - _LOGGER.warning("Got exception while fetching the state: %s", exc) - self._received_map_name_previously = False - finally: - counter = counter - 1 - return map_name - - def _handle_map_data(self, map_name: str): - _LOGGER.debug("Retrieving map from Xiaomi cloud") - store_map_path = self._store_map_path if self._store_map_raw else None - map_data, map_stored = self._device.get_map(map_name, self._colors, self._drawables, self._texts, - self._sizes, self._image_config, store_map_path) - if map_data is not None: - # noinspection PyBroadException - try: - _LOGGER.debug("Map data retrieved") - self._map_saved = map_stored - if map_data.image.is_empty: - _LOGGER.debug("Map is empty") - self._status = CameraStatus.EMPTY_MAP - if self._map_data is None or self._map_data.image.is_empty: - self._set_map_data(map_data) - else: - _LOGGER.debug("Map is ok") - self._set_map_data(map_data) - self._status = CameraStatus.OK - except: - _LOGGER.warning("Unable to parse map data") - self._status = CameraStatus.UNABLE_TO_PARSE_MAP - else: - self._logged_in = False - _LOGGER.warning("Unable to retrieve map data") - self._status = CameraStatus.UNABLE_TO_RETRIEVE_MAP - - def _set_map_data(self, map_data: MapData): - img_byte_arr = io.BytesIO() - map_data.image.data.save(img_byte_arr, format='PNG') - self._image = img_byte_arr.getvalue() - self._map_data = map_data - self._store_image() - - def _create_device(self, user_id: str, device_id: str, model: str) -> XiaomiCloudVacuum: - self._used_api = self._detect_api(model) - if self._used_api == CONF_AVAILABLE_API_XIAOMI: - return XiaomiVacuum(self._connector, self._country, user_id, device_id, model) - if self._used_api == CONF_AVAILABLE_API_VIOMI: - return ViomiVacuum(self._connector, self._country, user_id, device_id, model) - if self._used_api == CONF_AVAILABLE_API_ROIDMI: - return RoidmiVacuum(self._connector, self._country, user_id, device_id, model) - if self._used_api == CONF_AVAILABLE_API_DREAME: - return DreameVacuum(self._connector, self._country, user_id, device_id, model) - return UnsupportedVacuum(self._connector, self._country, user_id, device_id, model) - - def _detect_api(self, model: str) -> Optional[str]: - if self._forced_api is not None: - return self._forced_api - if model in API_EXCEPTIONS: - return API_EXCEPTIONS[model] - - def list_contains_model(prefixes, model_to_check): - return len(list(filter(lambda x: model_to_check.startswith(x), prefixes))) > 0 - - filtered = list(filter(lambda x: list_contains_model(x[1], model), AVAILABLE_APIS.items())) - if len(filtered) > 0: - return filtered[0][0] - return None - - def _store_image(self): - if self._store_map_image: - try: - image = Image.open(io.BytesIO(self._image)) - image.save(f"{self._store_map_path}/map_image_{self._device.model}.png") - except: - _LOGGER.warning("Error while saving image") - - -class CameraStatus(Enum): - EMPTY_MAP = 'Empty map' - FAILED_LOGIN = 'Failed to login' - FAILED_TO_RETRIEVE_DEVICE = 'Failed to retrieve device' - FAILED_TO_RETRIEVE_MAP_FROM_VACUUM = 'Failed to retrieve map from vacuum' - INITIALIZING = 'Initializing' - NOT_LOGGED_IN = 'Not logged in' - OK = 'OK' - LOGGED_IN = 'Logged in' - TWO_FACTOR_AUTH_REQUIRED = 'Two factor auth required (see logs)' - UNABLE_TO_PARSE_MAP = 'Unable to parse map' - UNABLE_TO_RETRIEVE_MAP = 'Unable to retrieve map' - - def __str__(self): - return str(self._value_) + def should_poll(self): + return self._manager.should_poll diff --git a/custom_components/xiaomi_cloud_map_extractor/const.py b/custom_components/xiaomi_cloud_map_extractor/const.py index 5bdd5c6..0e6cb76 100644 --- a/custom_components/xiaomi_cloud_map_extractor/const.py +++ b/custom_components/xiaomi_cloud_map_extractor/const.py @@ -212,3 +212,23 @@ } MM = 50 + +DEFAULT_TRIMS = { + CONF_LEFT: 0, + CONF_RIGHT: 0, + CONF_TOP: 0, + CONF_BOTTOM: 0 +} + +DEFAULT_SIZES = { + CONF_SIZE_VACUUM_RADIUS: 6, + CONF_SIZE_PATH_WIDTH: 1, + CONF_SIZE_IGNORED_OBSTACLE_RADIUS: 3, + CONF_SIZE_IGNORED_OBSTACLE_WITH_PHOTO_RADIUS: 3, + CONF_SIZE_OBSTACLE_RADIUS: 3, + CONF_SIZE_OBSTACLE_WITH_PHOTO_RADIUS: 3, + CONF_SIZE_CHARGER_RADIUS: 6 +} + +DEFAULT_MAP_TRANSFORM = {CONF_SCALE: 1, CONF_ROTATE: 0, CONF_TRIM: DEFAULT_TRIMS} +DEFAULT_STORE_MAP_PATH = "/tmp" diff --git a/custom_components/xiaomi_cloud_map_extractor/enums.py b/custom_components/xiaomi_cloud_map_extractor/enums.py new file mode 100644 index 0000000..f4b5423 --- /dev/null +++ b/custom_components/xiaomi_cloud_map_extractor/enums.py @@ -0,0 +1,18 @@ +from enum import Enum + + +class CameraStatus(Enum): + EMPTY_MAP = 'Empty map' + FAILED_LOGIN = 'Failed to login' + FAILED_TO_RETRIEVE_DEVICE = 'Failed to retrieve device' + FAILED_TO_RETRIEVE_MAP_FROM_VACUUM = 'Failed to retrieve map from vacuum' + INITIALIZING = 'Initializing' + NOT_LOGGED_IN = 'Not logged in' + OK = 'OK' + LOGGED_IN = 'Logged in' + TWO_FACTOR_AUTH_REQUIRED = 'Two factor auth required (see logs)' + UNABLE_TO_PARSE_MAP = 'Unable to parse map' + UNABLE_TO_RETRIEVE_MAP = 'Unable to retrieve map' + + def __str__(self): + return str(self._value_) diff --git a/custom_components/xiaomi_cloud_map_extractor/vacuum_manager.py b/custom_components/xiaomi_cloud_map_extractor/vacuum_manager.py new file mode 100644 index 0000000..2ab2570 --- /dev/null +++ b/custom_components/xiaomi_cloud_map_extractor/vacuum_manager.py @@ -0,0 +1,338 @@ +import io +import logging +import time +from typing import List, Optional + +from custom_components.xiaomi_cloud_map_extractor.common.map_data import MapData +from custom_components.xiaomi_cloud_map_extractor.types import Colors, Drawables, ImageConfig, Sizes, Texts + +try: + from miio import RoborockVacuum, DeviceException +except ImportError: + from miio import Vacuum as RoborockVacuum, DeviceException + +import PIL.Image as Image + +from homeassistant.const import CONF_HOST, CONF_NAME, CONF_PASSWORD, CONF_TOKEN, CONF_USERNAME + +from custom_components.xiaomi_cloud_map_extractor.common.map_data_parser import MapDataParser +from custom_components.xiaomi_cloud_map_extractor.common.xiaomi_cloud_connector import XiaomiCloudConnector +from custom_components.xiaomi_cloud_map_extractor.const import * +from custom_components.xiaomi_cloud_map_extractor.dreame.vacuum import DreameVacuum +from custom_components.xiaomi_cloud_map_extractor.enums import CameraStatus +from custom_components.xiaomi_cloud_map_extractor.roidmi.vacuum import RoidmiVacuum +from custom_components.xiaomi_cloud_map_extractor.unsupported.vacuum import UnsupportedVacuum +from custom_components.xiaomi_cloud_map_extractor.viomi.vacuum import ViomiVacuum +from custom_components.xiaomi_cloud_map_extractor.xiaomi.vacuum import XiaomiVacuum + + +_LOGGER = logging.getLogger(__name__) + +DEVICE_MAPPING = { + CONF_AVAILABLE_API_XIAOMI: XiaomiVacuum, + CONF_AVAILABLE_API_VIOMI: ViomiVacuum, + CONF_AVAILABLE_API_ROIDMI: RoidmiVacuum, + CONF_AVAILABLE_API_DREAME: DreameVacuum, +} + +STATUS_LOG_LEVEL = { + CameraStatus.FAILED_TO_RETRIEVE_DEVICE: _LOGGER.error, + CameraStatus.UNABLE_TO_PARSE_MAP: _LOGGER.warning, + CameraStatus.UNABLE_TO_RETRIEVE_MAP: _LOGGER.warning +} + + +class VacuumManager: + def __init__(self, config): + + host: str = config[CONF_HOST] + token: str = config[CONF_TOKEN] + username: str = config[CONF_USERNAME] + password: str = config[CONF_PASSWORD] + + drawables = config.get(CONF_DRAW, []) + room_colors = config.get(CONF_ROOM_COLORS, {}) + colors: Colors = config.get(CONF_COLORS, {}) + + for room, color in room_colors.items(): + colors[f"{COLOR_ROOM_PREFIX}{room}"] = color + + self._vacuum = RoborockVacuum(host, token) + self._connector = XiaomiCloudConnector(username, password) + + self._name: str = config.get(CONF_NAME, DEFAULT_NAME) + self._should_poll: bool = config.get(CONF_AUTO_UPDATE, True) + self._image_config: ImageConfig = config.get(CONF_MAP_TRANSFORM, DEFAULT_MAP_TRANSFORM) + self._colors: Colors = colors + self._drawables: Drawables = CONF_AVAILABLE_DRAWABLES[1:] if DRAWABLE_ALL in drawables else drawables + self._sizes: Sizes = config.get(CONF_SIZES, DEFAULT_SIZES) + self._texts: Texts = config.get(CONF_TEXTS, []) + self._country: str = config.get(CONF_COUNTRY) + self._allowed_attributes: List[str] = config.get(CONF_ATTRIBUTES, []) + self._store_map_raw: bool = config.get(CONF_STORE_MAP_RAW, False) + self._store_map_image: bool = config.get(CONF_STORE_MAP_IMAGE) + self._store_map_path: str = config.get(CONF_STORE_MAP_PATH, DEFAULT_STORE_MAP_PATH) + self._forced_api: str = config.get(CONF_FORCE_API) + + self._device = None + self._used_api = None + self._map_saved = None + self._image = None + self._map_data = None + self._logged_in = False + + self._logged_in_previously = True + self._received_map_name_previously = True + + self._attributes = {} + + self._status = CameraStatus.INITIALIZING + + @property + def image(self) -> Optional[bytes]: + return self._image + + @property + def name(self): + return self._name + + @property + def attributes(self): + return self._attributes + + @property + def should_poll(self): + return self._should_poll + + def turn_on(self): + self._should_poll = True + + def turn_off(self): + self._should_poll = False + + def _get_attributes_data(self): + map_data = self._map_data + + rooms = [] + if self._map_data.rooms is not None: + rooms = dict( + filter(lambda x: x[0] is not None, map(lambda x: (x[0], x[1].name), self._map_data.rooms.items()))) + if len(rooms) == 0: + rooms = list(self._map_data.rooms.keys()) + + attributes = { + ATTRIBUTE_CALIBRATION: map_data.calibration(), + ATTRIBUTE_CHARGER: map_data.charger, + ATTRIBUTE_CLEANED_ROOMS: map_data.cleaned_rooms, + ATTRIBUTE_COUNTRY: self._country, + ATTRIBUTE_GOTO: map_data.goto, + ATTRIBUTE_GOTO_PATH: map_data.goto_path, + ATTRIBUTE_GOTO_PREDICTED_PATH: map_data.predicted_path, + ATTRIBUTE_IGNORED_OBSTACLES: map_data.ignored_obstacles, + ATTRIBUTE_IGNORED_OBSTACLES_WITH_PHOTO: map_data.ignored_obstacles_with_photo, + ATTRIBUTE_IMAGE: map_data.image, + ATTRIBUTE_IS_EMPTY: map_data.image.is_empty, + ATTRIBUTE_MAP_NAME: map_data.map_name, + ATTRIBUTE_NO_GO_AREAS: map_data.no_go_areas, + ATTRIBUTE_NO_MOPPING_AREAS: map_data.no_mopping_areas, + ATTRIBUTE_OBSTACLES: map_data.obstacles, + ATTRIBUTE_OBSTACLES_WITH_PHOTO: map_data.obstacles_with_photo, + ATTRIBUTE_PATH: map_data.path, + ATTRIBUTE_ROOM_NUMBERS: rooms, + ATTRIBUTE_ROOMS: map_data.rooms, + ATTRIBUTE_VACUUM_POSITION: map_data.vacuum_position, + ATTRIBUTE_VACUUM_ROOM: map_data.vacuum_room, + ATTRIBUTE_VACUUM_ROOM_NAME: map_data.vacuum_room_name, + ATTRIBUTE_WALLS: map_data.walls, + ATTRIBUTE_ZONES: map_data.zones + } + + return attributes + + def _update_attributes(self): + attributes = {} + if self._map_data is not None: + data = self._get_attributes_data() + + for name, value in data.items(): + if name in self._allowed_attributes: + attributes[name] = value + + if self._store_map_raw: + attributes[ATTRIBUTE_MAP_SAVED] = self._map_saved + + if self._device is not None: + attributes[ATTR_MODEL] = self._device.model + attributes[ATTR_USED_API] = self._used_api + + if self._connector.two_factor_auth_url is not None: + attributes[ATTR_TWO_FACTOR_AUTH] = self._connector.two_factor_auth_url + + self._attributes = attributes + + def update(self, now): + counter = 10 + + if self._status != CameraStatus.TWO_FACTOR_AUTH_REQUIRED and not self._logged_in: + self._handle_login() + + if self._device is None and self._logged_in: + self._handle_device() + + map_name = self._handle_map_name(counter) + if map_name == "retry" and self._device is not None: + self._set_status(CameraStatus.FAILED_TO_RETRIEVE_MAP_FROM_VACUUM) + + self._received_map_name_previously = map_name != "retry" + if self._logged_in and map_name != "retry" and self._device is not None: + self._handle_map_data(map_name) + + else: + exists = self._device is not None + + _LOGGER.debug( + f"Unable to retrieve map ({now}), " + f"Logged in: {self._logged_in} | " + f"Map name: {map_name} | " + f"Device retrieved: {exists}" + ) + + message = str(self._status) + map_data = MapDataParser.create_empty(self._colors, message) + self._set_map_data(map_data) + + self._logged_in_previously = self._logged_in + + self._update_attributes() + + def _handle_login(self): + _LOGGER.debug("Logging in...") + + self._logged_in = self._connector.login() + if self._logged_in is None: + self._set_status(CameraStatus.TWO_FACTOR_AUTH_REQUIRED) + + elif self._logged_in: + self._set_status(CameraStatus.LOGGED_IN) + + else: + self._set_status(CameraStatus.FAILED_LOGIN) + + if self._logged_in_previously: + _LOGGER.error("Unable to log in, check credentials") + + def _handle_device(self): + _LOGGER.debug(f"Retrieving device info, country: {self._country}") + + country, user_id, device_id, model = self._connector.get_device_details(self._vacuum.token, self._country) + if model is not None: + self._country = country + _LOGGER.debug(f"Retrieved device model: {model}") + + self._used_api = self._detect_api(model) + + device_init = DEVICE_MAPPING.get(self._used_api, UnsupportedVacuum) + self._device = device_init(self._connector, self._country, user_id, device_id, model) + _LOGGER.debug(f"Created device, used api: {self._used_api}") + + else: + self._set_status(CameraStatus.FAILED_TO_RETRIEVE_DEVICE) + + def _handle_map_name(self, counter): + map_name = "retry" + if self._device is not None and not self._device.should_get_map_from_vacuum(): + map_name = "0" + + while map_name == "retry" and counter > 0: + _LOGGER.debug("Retrieving map name from device") + time.sleep(0.1) + + try: + map_name = self._vacuum.map()[0] + _LOGGER.debug("Map name %s", map_name) + + except OSError as exc: + _LOGGER.error(f"Got OSError while fetching the state: {str(exc)}") + + except DeviceException as exc: + if self._received_map_name_previously: + _LOGGER.warning(f"Got exception while fetching the state: {str(exc)}") + + self._received_map_name_previously = False + + finally: + counter = counter - 1 + + return map_name + + def _handle_map_data(self, map_name): + _LOGGER.debug("Retrieving map from Xiaomi cloud") + store_map_path = self._store_map_path if self._store_map_raw else None + map_data, map_stored = self._device.get_map(map_name, self._colors, self._drawables, self._texts, + self._sizes, self._image_config, store_map_path) + if map_data is not None: + # noinspection PyBroadException + try: + _LOGGER.debug("Map data retrieved") + self._set_map_data(map_data) + self._map_saved = map_stored + + if self._map_data.image.is_empty: + self._set_status(CameraStatus.EMPTY_MAP) + + if self._map_data is None or self._map_data.image.is_empty: + self._set_map_data(map_data) + + else: + self._set_map_data(map_data) + + self._set_status(CameraStatus.OK) + + except Exception as ex: + self._set_status(CameraStatus.UNABLE_TO_PARSE_MAP, ex) + + else: + self._logged_in = False + self._set_status(CameraStatus.UNABLE_TO_RETRIEVE_MAP) + + def _set_status(self, status, ex: Optional[Exception] = None): + log = STATUS_LOG_LEVEL.get(status, _LOGGER.debug) + + log_message = status + if ex is not None: + log_message = f"{status}, Error: {str(ex)}" + + self._status = status + log(log_message) + + def _set_map_data(self, map_data: MapData): + img_byte_arr = io.BytesIO() + map_data.image.data.save(img_byte_arr, format='PNG') + self._image = img_byte_arr.getvalue() + self._map_data = map_data + self._store_image() + + def _detect_api(self, model: str): + if self._forced_api is not None: + return self._forced_api + if model in API_EXCEPTIONS: + return API_EXCEPTIONS[model] + + def list_contains_model(prefixes): + return len(list(filter(lambda x: model.startswith(x), prefixes))) > 0 + + filtered = list(filter(lambda x: list_contains_model(x[1]), AVAILABLE_APIS.items())) + if len(filtered) > 0: + return filtered[0][0] + + return CONF_AVAILABLE_API_XIAOMI + + def _store_image(self): + if self._store_map_image: + try: + if self._image is not None: + image = Image.open(io.BytesIO(self._image)) + image.save(f"{self._store_map_path}/map_image_{self._device.model}.png") + + except Exception as ex: + _LOGGER.warning(f"Error while saving image, Error: {str(ex)}")