Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions roborock/device_features.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]})
Expand Down
1 change: 1 addition & 0 deletions roborock/devices/device_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
52 changes: 52 additions & 0 deletions roborock/devices/rpc/v1_channel.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
ResponseMessage,
SecurityData,
V1RpcChannel,
create_blob_response_decoder,
create_map_response_decoder,
create_security_data,
decode_data_protocol_message,
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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:
Expand Down
17 changes: 17 additions & 0 deletions roborock/devices/traits/v1/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@
map_content,
maps,
network_info,
obstacle_photos,
rooms,
routines,
smart_wash_params,
Expand All @@ -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
Expand All @@ -131,6 +133,7 @@
"map_content",
"maps",
"network_info",
"obstacle_photos",
"rooms",
"routines",
"smart_wash_params",
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -383,6 +399,7 @@ def create(
rpc_channel,
mqtt_rpc_channel,
map_rpc_channel,
blob_rpc_channel,
add_dps_listener,
web_api,
device_cache,
Expand Down
2 changes: 1 addition & 1 deletion roborock/devices/traits/v1/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
_LOGGER = logging.getLogger(__name__)


V1ResponseData = dict | list | int | str
V1ResponseData = dict | list | int | str | bytes


class V1TraitDataConverter(ABC):
Expand Down
104 changes: 104 additions & 0 deletions roborock/devices/traits/v1/obstacle_photos.py
Original file line number Diff line number Diff line change
@@ -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
45 changes: 45 additions & 0 deletions roborock/protocols/v1_protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -24,6 +26,7 @@
__all__ = [
"SecurityData",
"create_security_data",
"create_blob_response_decoder",
"decode_data_protocol_message",
"decode_rpc_response",
"V1RpcChannel",
Expand Down Expand Up @@ -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."""

Expand Down
1 change: 1 addition & 0 deletions roborock/roborock_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,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"
Expand Down
Loading
Loading