roborock.devices.traits.v1.obstacle_photos

  1"""Trait for fetching obstacle photos from V1 vacuums."""
  2
  3from dataclasses import dataclass
  4
  5from roborock.data import RoborockBase
  6from roborock.devices.traits.v1 import common
  7from roborock.exceptions import RoborockException
  8from roborock.protocols.v1_protocol import V1RpcChannel
  9from roborock.roborock_typing import RoborockCommand
 10
 11_PHOTO_TYPE_SMALL = 1
 12_PHOTO_DATA_BLOCK_TYPE = 3
 13_MAP_OBJECT_PHOTO_ENABLED_BIT = 10
 14_TYPE_SIZE = 2
 15_HEADER_SIZE_SIZE = 2
 16_PAYLOAD_SIZE_SIZE = 4
 17_MIN_BLOCK_HEADER_SIZE = _TYPE_SIZE + _HEADER_SIZE_SIZE + _PAYLOAD_SIZE_SIZE
 18_IMAGE_HEADERS = (b"\x89PNG\r\n\x1a\n", b"\xff\xd8\xff")
 19
 20
 21@dataclass
 22class ObstaclePhoto(RoborockBase):
 23    """Obstacle photo content."""
 24
 25    photo_id: str
 26    image_content: bytes
 27
 28
 29class ObstaclePhotoConverter(common.V1TraitDataConverter):
 30    """Convert a decrypted get_photo payload to an obstacle photo."""
 31
 32    def convert(self, response: common.V1ResponseData) -> ObstaclePhoto:
 33        """Parse the response from the device into an obstacle photo."""
 34        if not isinstance(response, bytes):
 35            raise ValueError(f"Unexpected ObstaclePhotoTrait response format: {type(response)}")
 36        return ObstaclePhoto(photo_id="", image_content=parse_photo_data(response))
 37
 38
 39def parse_photo_data(response: bytes) -> bytes:
 40    """Parse the get_photo response payload and return image bytes.
 41
 42    Roborock's app parses get_photo as a sequence of little-endian typed blocks.
 43    Block type 3 contains the image bytes.
 44    """
 45    offset = 0
 46    while offset + _MIN_BLOCK_HEADER_SIZE <= len(response):
 47        block_type = int.from_bytes(response[offset : offset + _TYPE_SIZE], "little")
 48        header_size = int.from_bytes(
 49            response[offset + _TYPE_SIZE : offset + _TYPE_SIZE + _HEADER_SIZE_SIZE],
 50            "little",
 51        )
 52        payload_size = int.from_bytes(
 53            response[offset + _TYPE_SIZE + _HEADER_SIZE_SIZE : offset + _MIN_BLOCK_HEADER_SIZE],
 54            "little",
 55        )
 56        next_offset = offset + header_size + payload_size
 57        if header_size < _MIN_BLOCK_HEADER_SIZE or next_offset > len(response):
 58            raise RoborockException("Invalid obstacle photo payload")
 59
 60        if block_type == _PHOTO_DATA_BLOCK_TYPE:
 61            image_content = response[offset + header_size : next_offset]
 62            if not image_content.startswith(_IMAGE_HEADERS):
 63                raise RoborockException("Obstacle photo payload is not a supported image")
 64            return image_content
 65
 66        offset = next_offset
 67
 68    raise RoborockException("Obstacle photo payload does not contain photo data")
 69
 70
 71class ObstaclePhotoTrait(RoborockBase, common.V1TraitMixin):
 72    """Trait for fetching obstacle photos."""
 73
 74    command = RoborockCommand.GET_PHOTO
 75    converter = ObstaclePhotoConverter()
 76    blob_rpc_channel = True
 77    requires_feature = "is_ai_recognition_obstacle_supported"
 78
 79    def __init__(self, standard_rpc_channel: V1RpcChannel) -> None:
 80        """Initialize the obstacle photo trait."""
 81        super().__init__()
 82        self._standard_rpc_channel = standard_rpc_channel
 83
 84    async def get_enabled(self) -> bool:
 85        """Return whether map object photo capture is enabled on the vacuum."""
 86        response = await self._standard_rpc_channel.send_command(RoborockCommand.GET_CAMERA_STATUS)
 87        if not isinstance(response, list) or not response or not isinstance(response[0], int):
 88            raise RoborockException("get_camera_status response did not contain camera status")
 89        return bool((response[0] >> _MAP_OBJECT_PHOTO_ENABLED_BIT) & 1)
 90
 91    async def get_photo(self, photo_id: str) -> ObstaclePhoto:
 92        """Fetch the small obstacle photo for a map photo ID.
 93
 94        Photo IDs are available as ``Obstacle.details.photo_name`` in
 95        ``ParsedMapData.map_data.obstacles_with_photo`` and
 96        ``ParsedMapData.map_data.ignored_obstacles_with_photo``.
 97        """
 98        response = await self.rpc_channel.send_command(
 99            self.command,
100            params={"img_id": photo_id, "type": _PHOTO_TYPE_SMALL},
101        )
102        photo = self.converter.convert(response)
103        photo.photo_id = photo_id
104        return photo