chore: clean up security location

Author: Lash-L

roborock/devices/device_manager.py

@@ -252,7 +252,6 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                     device_cache=device_cache,
                     map_parser_config=map_parser_config,
                     region=user_data.region,
-                    security_data=channel.security_data,
                 )
             case DeviceVersion.A01:
                 channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)

roborock/devices/rpc/v1_channel.py

@@ -162,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.
 
@@ -214,11 +258,6 @@ def is_mqtt_connected(self) -> bool:
         """
         return self._mqtt_channel.is_connected and self._mqtt_unsub is not None
 
-    @property
-    def security_data(self) -> SecurityData:
-        """Return security data used for map/blob RPC commands."""
-        return self._security_data
-
     @property
     def rpc_channel(self) -> V1RpcChannel:
         """Return the combined RPC channel that prefers local with a fallback to MQTT.
@@ -255,7 +294,8 @@ def map_rpc_channel(self) -> V1RpcChannel:
     def blob_rpc_channel(self) -> V1RpcChannel:
         """Return the blob RPC channel used for fetching binary content."""
         decoder = create_blob_response_decoder()
-        return RpcChannel(lambda: [self._create_mqtt_rpc_strategy(decoder)], self._logger)
+        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."""

roborock/devices/traits/v1/__init__.py

@@ -64,7 +64,7 @@
 from roborock.devices.traits import Trait
 from roborock.exceptions import RoborockException
 from roborock.map.map_parser import MapParserConfig
-from roborock.protocols.v1_protocol import SecurityData, V1RpcChannel, decode_data_protocol_message
+from roborock.protocols.v1_protocol import V1RpcChannel, decode_data_protocol_message
 from roborock.roborock_message import RoborockDataProtocol, RoborockMessage
 from roborock.web_api import UserWebApiClient
 
@@ -190,15 +190,13 @@ def __init__(
         device_cache: DeviceCache,
         map_parser_config: MapParserConfig | None = None,
         region: str | None = None,
-        security_data: SecurityData | None = None,
     ) -> None:
         """Initialize the V1TraitProps."""
         self._device_uid = device_uid
         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._security_data = security_data
         self._web_api = web_api
         self._device_cache = device_cache
         self._region = region
@@ -282,12 +280,8 @@ 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._security_data is not None
-            and self._is_supported(ObstaclePhotoTrait, "obstacle_photos", dock_features)
-        ):
-            obstacle_photos = ObstaclePhotoTrait(self._rpc_channel, self._security_data)
+        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
 
@@ -396,7 +390,6 @@ def create(
     device_cache: DeviceCache,
     map_parser_config: MapParserConfig | None = None,
     region: str | None = None,
-    security_data: SecurityData | None = None,
 ) -> PropertiesApi:
     """Create traits for V1 devices."""
     return PropertiesApi(
@@ -412,5 +405,4 @@ def create(
         device_cache,
         map_parser_config,
         region=region,
-        security_data=security_data,
     )

roborock/devices/traits/v1/obstacle_photos.py

@@ -5,7 +5,7 @@
 from roborock.data import RoborockBase
 from roborock.devices.traits.v1 import common
 from roborock.exceptions import RoborockException
-from roborock.protocols.v1_protocol import SecurityData, V1RpcChannel
+from roborock.protocols.v1_protocol import V1RpcChannel
 from roborock.roborock_typing import RoborockCommand
 
 _PHOTO_TYPE_SMALL = 1
@@ -76,11 +76,10 @@ class ObstaclePhotoTrait(RoborockBase, common.V1TraitMixin):
     blob_rpc_channel = True
     requires_feature = "is_ai_recognition_obstacle_supported"
 
-    def __init__(self, standard_rpc_channel: V1RpcChannel, security_data: SecurityData) -> None:
+    def __init__(self, standard_rpc_channel: V1RpcChannel) -> None:
         """Initialize the obstacle photo trait."""
         super().__init__()
         self._standard_rpc_channel = standard_rpc_channel
-        self._security_data = security_data
 
     async def get_enabled(self) -> bool:
         """Return whether map object photo capture is enabled on the vacuum."""
@@ -91,21 +90,9 @@ async def get_enabled(self) -> bool:
 
     async def get_photo(self, photo_id: str, photo_type: int = _PHOTO_TYPE_SMALL) -> ObstaclePhoto:
         """Fetch an obstacle photo by its map photo id."""
-        public_key = await self._standard_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"]
         response = await self.rpc_channel.send_command(
             self.command,
-            params={
-                "security": {
-                    "pub_key": public_key["pub_key"],
-                    "cipher_suite": 0,
-                },
-                "endpoint": security["endpoint"],
-                "nonce": security["nonce"],
-                "data_filter": {"img_id": photo_id, "type": photo_type},
-            },
+            params={"img_id": photo_id, "type": photo_type},
         )
         photo = self.converter.convert(response)
         photo.photo_id = photo_id

tests/devices/rpc/test_v1_channel.py

@@ -4,6 +4,7 @@
 and failure modes, ensuring the V1Channel behaves correctly in various scenarios.
 """
 
+import gzip
 import json
 import logging
 from collections.abc import Iterator
@@ -163,6 +164,12 @@ def setup_map_rpc_channel(v1_channel: V1Channel) -> V1RpcChannel:
     return v1_channel.map_rpc_channel
 
 
+@pytest.fixture(name="blob_rpc_channel")
+def setup_blob_rpc_channel(v1_channel: V1Channel) -> V1RpcChannel:
+    """Fixture to set up the Blob RPC channel for tests."""
+    return v1_channel.blob_rpc_channel
+
+
 @pytest.fixture(name="warning_caplog")
 def setup_warning_caplog(caplog: pytest.LogCaptureFixture) -> pytest.LogCaptureFixture:
     """Fixture to capture warning messages."""
@@ -632,6 +639,75 @@ async def test_v1_channel_send_map_command(
     assert result == decompressed_map_data
 
 
+async def test_v1_channel_send_blob_command(
+    blob_rpc_channel: V1RpcChannel,
+    mock_mqtt_channel: FakeChannel,
+) -> None:
+    """Test that the blob channel adds security parameters and decodes the response."""
+    public_key = {"n": "abc", "e": "010001"}
+    mock_mqtt_channel.response_queue.append(
+        RoborockMessage(
+            protocol=RoborockMessageProtocol.RPC_RESPONSE,
+            payload=json.dumps({"dps": {"102": json.dumps({"id": 12345, "result": {"pub_key": public_key}})}}).encode(),
+        )
+    )
+
+    blob_data = b"blob response"
+    compressed = gzip.compress(blob_data)
+    header_size = 24
+    payload = bytearray(header_size + len(compressed))
+    payload[:8] = b"ROBOROCK"
+    payload[8:12] = (12346).to_bytes(4, "little")
+    payload[16:18] = header_size.to_bytes(2, "little")
+    payload[20:24] = len(compressed).to_bytes(4, "little")
+    payload[header_size:] = compressed
+    mock_mqtt_channel.response_queue.append(
+        RoborockMessage(protocol=RoborockMessageProtocol.MAP_RESPONSE, payload=bytes(payload))
+    )
+
+    result = await blob_rpc_channel.send_command(
+        RoborockCommand.GET_PHOTO,
+        params={"img_id": "photo-id", "type": 1},
+    )
+
+    assert result == blob_data
+    sent_payload = mock_mqtt_channel.published_messages[-1].payload
+    assert sent_payload is not None
+    request_payload = json.loads(sent_payload)
+    request = json.loads(request_payload["dps"]["101"])
+    assert request["method"] == RoborockCommand.GET_PHOTO
+    assert request["params"] == {
+        "security": {
+            "pub_key": public_key,
+            "cipher_suite": 0,
+        },
+        "endpoint": TEST_SECURITY_DATA.endpoint,
+        "nonce": TEST_SECURITY_DATA.nonce.hex(),
+        "data_filter": {"img_id": "photo-id", "type": 1},
+    }
+
+
+async def test_v1_channel_send_blob_command_rejects_invalid_public_key(
+    blob_rpc_channel: V1RpcChannel,
+    mock_mqtt_channel: FakeChannel,
+) -> None:
+    """Test that the blob channel rejects an invalid public key response."""
+    mock_mqtt_channel.response_queue.append(
+        RoborockMessage(
+            protocol=RoborockMessageProtocol.RPC_RESPONSE,
+            payload=json.dumps({"dps": {"102": json.dumps({"id": 12345, "result": {}})}}).encode(),
+        )
+    )
+
+    with pytest.raises(RoborockException, match="did not contain a public key"):
+        await blob_rpc_channel.send_command(
+            RoborockCommand.GET_PHOTO,
+            params={"img_id": "photo-id", "type": 1},
+        )
+
+    assert len(mock_mqtt_channel.published_messages) == 1
+
+
 async def test_v1_channel_add_dps_listener(
     v1_channel: V1Channel,
     mock_mqtt_channel: FakeChannel,

tests/devices/traits/v1/test_obstacle_photos.py

@@ -11,7 +11,7 @@
     parse_photo_data,
 )
 from roborock.exceptions import RoborockException
-from roborock.protocols.v1_protocol import SecurityData, create_blob_response_decoder
+from roborock.protocols.v1_protocol import create_blob_response_decoder
 from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
 from roborock.roborock_typing import RoborockCommand
 
@@ -33,10 +33,7 @@ def _block(block_type: int, payload: bytes, header_size: int = 8) -> bytes:
 def obstacle_photo_trait(device: RoborockDevice, mock_blob_rpc_channel: AsyncMock) -> ObstaclePhotoTrait:
     """Create an ObstaclePhotoTrait instance with mocked dependencies."""
     assert device.v1_properties
-    trait = ObstaclePhotoTrait(
-        device.v1_properties.status.rpc_channel,
-        SecurityData(endpoint="endpoint", nonce=b"1234567890abcdef"),
-    )
+    trait = ObstaclePhotoTrait(device.v1_properties.status.rpc_channel)
     trait._rpc_channel = mock_blob_rpc_channel
     return trait
 
@@ -101,26 +98,15 @@ async def test_get_enabled(obstacle_photo_trait: ObstaclePhotoTrait, mock_rpc_ch
 async def test_get_photo(
     obstacle_photo_trait: ObstaclePhotoTrait,
     mock_blob_rpc_channel: AsyncMock,
-    mock_rpc_channel: AsyncMock,
 ) -> None:
     """Test fetching and parsing obstacle photo content."""
-    mock_rpc_channel.send_command.return_value = {"pub_key": {"n": "abc", "e": "010001"}}
     mock_blob_rpc_channel.send_command.return_value = _block(3, PNG_BYTES)
 
     photo = await obstacle_photo_trait.get_photo("photo-id")
 
     assert photo.photo_id == "photo-id"
     assert photo.image_content == PNG_BYTES
-    mock_rpc_channel.send_command.assert_called_once_with(RoborockCommand.GET_RANDOM_PKEY)
     mock_blob_rpc_channel.send_command.assert_called_once_with(
         RoborockCommand.GET_PHOTO,
-        params={
-            "security": {
-                "pub_key": {"n": "abc", "e": "010001"},
-                "cipher_suite": 0,
-            },
-            "endpoint": "endpoint",
-            "nonce": "31323334353637383930616263646566",
-            "data_filter": {"img_id": "photo-id", "type": 1},
-        },
+        params={"img_id": "photo-id", "type": 1},
     )