feat: add photo for obstacle avoidance (#880)

* feat: add photo for obstacle avoidance

* chore: clean up security location

* chore: mild clean ups

Author: Lash-L

roborock/device_features.py

@@ -539,6 +539,12 @@ class DeviceFeatures(RoborockBase):
             ]
         }
     )
+    is_ai_recognition_setting_supported: bool = field(
+        metadata={"product_features": [ProductFeatures.AIRECOGNITION_SETTING]}
+    )
+    is_ai_recognition_obstacle_supported: bool = field(
+        metadata={"product_features": [ProductFeatures.AIRECOGNITION_OBSTACLE]}
+    )
     is_pure_clean_mop_supported: bool = field(metadata={"product_features": [ProductFeatures.CLEANMODE_PURECLEANMOP]})
     is_new_remote_view_supported: bool = field(metadata={"product_features": [ProductFeatures.REMOTE_BACK]})
     is_max_plus_mode_supported: bool = field(metadata={"product_features": [ProductFeatures.CLEANMODE_MAXPLUS]})

roborock/devices/device_manager.py

@@ -246,6 +246,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                     channel.rpc_channel,
                     channel.mqtt_rpc_channel,
                     channel.map_rpc_channel,
+                    channel.blob_rpc_channel,
                     channel.add_dps_listener,
                     web_api,
                     device_cache=device_cache,

roborock/devices/rpc/v1_channel.py

@@ -29,6 +29,7 @@
     ResponseMessage,
     SecurityData,
     V1RpcChannel,
+    create_blob_response_decoder,
     create_map_response_decoder,
     create_security_data,
     decode_data_protocol_message,
@@ -161,6 +162,50 @@ def find_response(response_message: RoborockMessage) -> None:
         return result
 
 
+class BlobRpcChannel(V1RpcChannel):
+    """RPC channel that adds the security envelope required for blob requests."""
+
+    def __init__(
+        self,
+        rpc_channel: V1RpcChannel,
+        raw_blob_rpc_channel: V1RpcChannel,
+        security_data: SecurityData,
+    ) -> None:
+        """Initialize the blob RPC channel."""
+        self._rpc_channel = rpc_channel
+        self._raw_blob_rpc_channel = raw_blob_rpc_channel
+        self._security_data = security_data
+
+    async def send_command(
+        self,
+        method: CommandType,
+        *,
+        response_type: type[_T] | None = None,
+        params: ParamsType = None,
+    ) -> _T | Any:
+        """Send a blob command after adding the device security parameters."""
+        public_key = await self._rpc_channel.send_command(RoborockCommand.GET_RANDOM_PKEY)
+        if not isinstance(public_key, dict) or not isinstance(public_key.get("pub_key"), dict):
+            raise RoborockException("get_random_pkey response did not contain a public key")
+        security = self._security_data.to_dict()["security"]
+        blob_params = {
+            "security": {
+                "pub_key": public_key["pub_key"],
+                "cipher_suite": 0,
+            },
+            "endpoint": security["endpoint"],
+            "nonce": security["nonce"],
+            "data_filter": params,
+        }
+        if response_type is not None:
+            return await self._raw_blob_rpc_channel.send_command(
+                method,
+                response_type=response_type,
+                params=blob_params,
+            )
+        return await self._raw_blob_rpc_channel.send_command(method, params=blob_params)
+
+
 class V1Channel(Channel):
     """Unified V1 protocol channel with automatic MQTT/local connection handling.
 
@@ -245,6 +290,13 @@ def map_rpc_channel(self) -> V1RpcChannel:
         decoder = create_map_response_decoder(security_data=self._security_data)
         return RpcChannel(lambda: [self._create_mqtt_rpc_strategy(decoder)], self._logger)
 
+    @property
+    def blob_rpc_channel(self) -> V1RpcChannel:
+        """Return the blob RPC channel used for fetching binary content."""
+        decoder = create_blob_response_decoder()
+        raw_blob_rpc_channel = RpcChannel(lambda: [self._create_mqtt_rpc_strategy(decoder)], self._logger)
+        return BlobRpcChannel(self.rpc_channel, raw_blob_rpc_channel, self._security_data)
+
     def _create_local_rpc_strategy(self) -> RpcStrategy | None:
         """Create the RPC strategy for local transport."""
         if self._local_channel is None or not self.is_local_connected:

roborock/devices/traits/v1/__init__.py

@@ -83,6 +83,7 @@
     map_content,
     maps,
     network_info,
+    obstacle_photos,
     rooms,
     routines,
     smart_wash_params,
@@ -105,6 +106,7 @@
 from .map_content import MapContentTrait
 from .maps import MapsTrait
 from .network_info import NetworkInfoTrait
+from .obstacle_photos import ObstaclePhotoTrait
 from .rooms import RoomsTrait
 from .routines import RoutinesTrait
 from .smart_wash_params import SmartWashParamsTrait
@@ -131,6 +133,7 @@
     "map_content",
     "maps",
     "network_info",
+    "obstacle_photos",
     "rooms",
     "routines",
     "smart_wash_params",
@@ -171,6 +174,7 @@ class PropertiesApi(Trait):
     dust_collection_mode: DustCollectionModeTrait | None = None
     wash_towel_mode: WashTowelModeTrait | None = None
     smart_wash_params: SmartWashParamsTrait | None = None
+    obstacle_photos: ObstaclePhotoTrait | None = None
 
     def __init__(
         self,
@@ -180,6 +184,7 @@ def __init__(
         rpc_channel: V1RpcChannel,
         mqtt_rpc_channel: V1RpcChannel,
         map_rpc_channel: V1RpcChannel,
+        blob_rpc_channel: V1RpcChannel,
         add_dps_listener: Callable[[Callable[[dict[RoborockDataProtocol, Any]], None]], Callable[[], None]],
         web_api: UserWebApiClient,
         device_cache: DeviceCache,
@@ -191,6 +196,7 @@ def __init__(
         self._rpc_channel = rpc_channel
         self._mqtt_rpc_channel = mqtt_rpc_channel
         self._map_rpc_channel = map_rpc_channel
+        self._blob_rpc_channel = blob_rpc_channel
         self._web_api = web_api
         self._device_cache = device_cache
         self._region = region
@@ -229,6 +235,8 @@ def _get_rpc_channel(self, trait: V1TraitMixin) -> V1RpcChannel:
         # to use the mqtt_rpc_channel (cloud only) instead of the rpc_channel (adaptive)
         if hasattr(trait, "mqtt_rpc_channel"):
             return self._mqtt_rpc_channel
+        elif hasattr(trait, "blob_rpc_channel"):
+            return self._blob_rpc_channel
         elif hasattr(trait, "map_rpc_channel"):
             return self._map_rpc_channel
         else:
@@ -272,6 +280,11 @@ async def discover_features(self) -> None:
             wash_towel_mode._rpc_channel = self._get_rpc_channel(wash_towel_mode)  # type: ignore[assignment]
             self.wash_towel_mode = wash_towel_mode
 
+        if self.obstacle_photos is None and self._is_supported(ObstaclePhotoTrait, "obstacle_photos", dock_features):
+            obstacle_photos = ObstaclePhotoTrait(self._rpc_channel)
+            obstacle_photos._rpc_channel = self._get_rpc_channel(obstacle_photos)
+            self.obstacle_photos = obstacle_photos
+
         # Dynamically create any traits that need to be populated
         for item in fields(self):
             if (trait := getattr(self, item.name, None)) is not None:
@@ -283,6 +296,8 @@ async def discover_features(self) -> None:
 
             # Union args may not be in declared order
             item_type = union_args[0] if union_args[1] is type(None) else union_args[1]
+            if item_type is ObstaclePhotoTrait:
+                continue
             if not self._is_supported(item_type, item.name, dock_features):
                 _LOGGER.debug("Trait '%s' not supported, skipping", item.name)
                 continue
@@ -369,6 +384,7 @@ def create(
     rpc_channel: V1RpcChannel,
     mqtt_rpc_channel: V1RpcChannel,
     map_rpc_channel: V1RpcChannel,
+    blob_rpc_channel: V1RpcChannel,
     add_dps_listener: Callable[[Callable[[dict[RoborockDataProtocol, Any]], None]], Callable[[], None]],
     web_api: UserWebApiClient,
     device_cache: DeviceCache,
@@ -383,6 +399,7 @@ def create(
         rpc_channel,
         mqtt_rpc_channel,
         map_rpc_channel,
+        blob_rpc_channel,
         add_dps_listener,
         web_api,
         device_cache,

roborock/devices/traits/v1/common.py

@@ -16,7 +16,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 
-V1ResponseData = dict | list | int | str
+V1ResponseData = dict | list | int | str | bytes
 
 
 class V1TraitDataConverter(ABC):

roborock/devices/traits/v1/obstacle_photos.py

@@ -0,0 +1,104 @@
+"""Trait for fetching obstacle photos from V1 vacuums."""
+
+from dataclasses import dataclass
+
+from roborock.data import RoborockBase
+from roborock.devices.traits.v1 import common
+from roborock.exceptions import RoborockException
+from roborock.protocols.v1_protocol import V1RpcChannel
+from roborock.roborock_typing import RoborockCommand
+
+_PHOTO_TYPE_SMALL = 1
+_PHOTO_DATA_BLOCK_TYPE = 3
+_MAP_OBJECT_PHOTO_ENABLED_BIT = 10
+_TYPE_SIZE = 2
+_HEADER_SIZE_SIZE = 2
+_PAYLOAD_SIZE_SIZE = 4
+_MIN_BLOCK_HEADER_SIZE = _TYPE_SIZE + _HEADER_SIZE_SIZE + _PAYLOAD_SIZE_SIZE
+_IMAGE_HEADERS = (b"\x89PNG\r\n\x1a\n", b"\xff\xd8\xff")
+
+
+@dataclass
+class ObstaclePhoto(RoborockBase):
+    """Obstacle photo content."""
+
+    photo_id: str
+    image_content: bytes
+
+
+class ObstaclePhotoConverter(common.V1TraitDataConverter):
+    """Convert a decrypted get_photo payload to an obstacle photo."""
+
+    def convert(self, response: common.V1ResponseData) -> ObstaclePhoto:
+        """Parse the response from the device into an obstacle photo."""
+        if not isinstance(response, bytes):
+            raise ValueError(f"Unexpected ObstaclePhotoTrait response format: {type(response)}")
+        return ObstaclePhoto(photo_id="", image_content=parse_photo_data(response))
+
+
+def parse_photo_data(response: bytes) -> bytes:
+    """Parse the get_photo response payload and return image bytes.
+
+    Roborock's app parses get_photo as a sequence of little-endian typed blocks.
+    Block type 3 contains the image bytes.
+    """
+    offset = 0
+    while offset + _MIN_BLOCK_HEADER_SIZE <= len(response):
+        block_type = int.from_bytes(response[offset : offset + _TYPE_SIZE], "little")
+        header_size = int.from_bytes(
+            response[offset + _TYPE_SIZE : offset + _TYPE_SIZE + _HEADER_SIZE_SIZE],
+            "little",
+        )
+        payload_size = int.from_bytes(
+            response[offset + _TYPE_SIZE + _HEADER_SIZE_SIZE : offset + _MIN_BLOCK_HEADER_SIZE],
+            "little",
+        )
+        next_offset = offset + header_size + payload_size
+        if header_size < _MIN_BLOCK_HEADER_SIZE or next_offset > len(response):
+            raise RoborockException("Invalid obstacle photo payload")
+
+        if block_type == _PHOTO_DATA_BLOCK_TYPE:
+            image_content = response[offset + header_size : next_offset]
+            if not image_content.startswith(_IMAGE_HEADERS):
+                raise RoborockException("Obstacle photo payload is not a supported image")
+            return image_content
+
+        offset = next_offset
+
+    raise RoborockException("Obstacle photo payload does not contain photo data")
+
+
+class ObstaclePhotoTrait(RoborockBase, common.V1TraitMixin):
+    """Trait for fetching obstacle photos."""
+
+    command = RoborockCommand.GET_PHOTO
+    converter = ObstaclePhotoConverter()
+    blob_rpc_channel = True
+    requires_feature = "is_ai_recognition_obstacle_supported"
+
+    def __init__(self, standard_rpc_channel: V1RpcChannel) -> None:
+        """Initialize the obstacle photo trait."""
+        super().__init__()
+        self._standard_rpc_channel = standard_rpc_channel
+
+    async def get_enabled(self) -> bool:
+        """Return whether map object photo capture is enabled on the vacuum."""
+        response = await self._standard_rpc_channel.send_command(RoborockCommand.GET_CAMERA_STATUS)
+        if not isinstance(response, list) or not response or not isinstance(response[0], int):
+            raise RoborockException("get_camera_status response did not contain camera status")
+        return bool((response[0] >> _MAP_OBJECT_PHOTO_ENABLED_BIT) & 1)
+
+    async def get_photo(self, photo_id: str) -> ObstaclePhoto:
+        """Fetch the small obstacle photo for a map photo ID.
+
+        Photo IDs are available as ``Obstacle.details.photo_name`` in
+        ``ParsedMapData.map_data.obstacles_with_photo`` and
+        ``ParsedMapData.map_data.ignored_obstacles_with_photo``.
+        """
+        response = await self.rpc_channel.send_command(
+            self.command,
+            params={"img_id": photo_id, "type": _PHOTO_TYPE_SMALL},
+        )
+        photo = self.converter.convert(response)
+        photo.photo_id = photo_id
+        return photo

roborock/protocols/v1_protocol.py

@@ -3,10 +3,12 @@
 from __future__ import annotations
 
 import base64
+import gzip
 import json
 import logging
 import secrets
 import struct
+import zlib
 from collections.abc import Callable
 from dataclasses import dataclass, field
 from enum import StrEnum
@@ -24,6 +26,7 @@
 __all__ = [
     "SecurityData",
     "create_security_data",
+    "create_blob_response_decoder",
     "decode_data_protocol_message",
     "decode_rpc_response",
     "V1RpcChannel",
@@ -261,6 +264,48 @@ class MapResponse:
     """The map data, decrypted and decompressed."""
 
 
+@dataclass
+class BlobResponse:
+    """Data structure for V1 blob responses."""
+
+    request_id: int
+    """The request ID of the blob response."""
+
+    data: bytes
+    """The blob data, decompressed."""
+
+
+def create_blob_response_decoder() -> Callable[[RoborockMessage], BlobResponse | None]:
+    """Create a decoder for V1 blob response messages.
+
+    Obstacle photos are acknowledged through the normal RPC response with
+    ``["ok"]`` and delivered later as a protocol-301 blob frame. The frame starts
+    with ``ROBOROCK``, stores the RPC request id at bytes 8-11, the header size
+    at bytes 16-17, and the gzip payload length at bytes 20-23.
+    """
+
+    def _decode_blob_response(message: RoborockMessage) -> BlobResponse | None:
+        """Decode a V1 blob response message."""
+        if message.protocol != RoborockMessageProtocol.MAP_RESPONSE:
+            return None
+        payload = message.payload
+        if not payload or not payload.startswith(b"ROBOROCK") or len(payload) < 24:
+            return None
+        request_id = int.from_bytes(payload[8:12], "little")
+        header_size = int.from_bytes(payload[16:18], "little")
+        payload_size = int.from_bytes(payload[20:24], "little")
+        end_offset = header_size + payload_size
+        if header_size < 24 or end_offset > len(payload):
+            raise RoborockException("Invalid V1 blob response format")
+        try:
+            data = Utils.decompress(payload[header_size:end_offset])
+        except (gzip.BadGzipFile, EOFError, zlib.error) as err:
+            raise RoborockException("Failed to decode blob message payload") from err
+        return BlobResponse(request_id=request_id, data=data)
+
+    return _decode_blob_response
+
+
 def create_map_response_decoder(security_data: SecurityData) -> Callable[[RoborockMessage], MapResponse | None]:
     """Create a decoder for V1 map response messages."""
 

roborock/roborock_typing.py

@@ -117,6 +117,7 @@ class RoborockCommand(str, Enum):
     GET_NETWORK_INFO = "get_network_info"
     GET_OFFLINE_MAP_STATUS = "get_offline_map_status"
     GET_PERSIST = "get_persist_map"
+    GET_PHOTO = "get_photo"
     GET_PROP = "get_prop"
     GET_RANDOM_PKEY = "get_random_pkey"
     GET_RECOVER_MAP = "get_recover_map"