🦎 q7: simplify map content parsing

arduano · arduano · commit c4f12d214987 · 2026-03-15T12:33:30.000+11:00
diff --git a/roborock/devices/traits/b01/q7/__init__.py b/roborock/devices/traits/b01/q7/__init__.py
@@ -21,7 +21,6 @@
 from roborock.devices.rpc.b01_q7_channel import send_decoded_command
 from roborock.devices.traits import Trait
 from roborock.devices.transport.mqtt_channel import MqttChannel
-from roborock.exceptions import RoborockException
 from roborock.protocols.b01_q7_protocol import B01_Q7_DPS, CommandType, ParamsType, Q7RequestMessage
 from roborock.roborock_message import RoborockB01Props
 from roborock.roborock_typing import RoborockB01Q7Methods
@@ -58,11 +57,6 @@ def __init__(self, channel: MqttChannel, *, device: HomeDataDevice, product: Hom
         self._device = device
         self._product = product
 
-        if not device.sn:
-            raise RoborockException(
-                "B01/Q7 map parsing requires the device serial number (HomeDataDevice.sn), but it was missing"
-            )
-
         self.clean_summary = CleanSummaryTrait(channel)
         self.map = MapTrait(channel)
         self.map_content = MapContentTrait(
diff --git a/roborock/devices/traits/b01/q7/map_content.py b/roborock/devices/traits/b01/q7/map_content.py
@@ -33,6 +33,9 @@ class MapContent(RoborockBase):
     map_data: MapData | None = None
     """Parsed map data (metadata for points on the map)."""
 
+    room_names: dict[int, str] | None = None
+    """Observed room-id to room-name mapping for the current map, if available."""
+
     raw_api_response: bytes | None = None
     """Raw bytes of the map payload from the device.
 
@@ -55,8 +58,8 @@ def __init__(
         map_trait: MapTrait,
         *,
         local_key: str,
-        serial: str,
-        model: str,
+        serial: str | None,
+        model: str | None,
         map_parser_config: B01MapParserConfig | None = None,
     ) -> None:
         super().__init__()
@@ -72,6 +75,7 @@ async def refresh(self) -> None:
         parsed = self.parse_map_content(raw_payload)
         self.image_content = parsed.image_content
         self.map_data = parsed.map_data
+        self.room_names = parsed.room_names
         self.raw_api_response = parsed.raw_api_response
 
     def parse_map_content(self, response: bytes) -> MapContent:
@@ -80,6 +84,11 @@ def parse_map_content(self, response: bytes) -> MapContent:
         Exposed so callers can re-parse cached map payload bytes without
         performing I/O.
         """
+        if not self._serial or not self._model:
+            raise RoborockException(
+                "B01 map parsing requires device serial number and model metadata, but they were missing"
+            )
+
         try:
             parsed_data = self._map_parser.parse(
                 response,
@@ -93,8 +102,13 @@ def parse_map_content(self, response: bytes) -> MapContent:
         if parsed_data.image_content is None:
             raise RoborockException("Failed to render B01 map image")
 
+        room_names = None
+        if parsed_data.map_data is not None:
+            room_names = parsed_data.map_data.additional_parameters.get("room_names")
+
         return MapContent(
             image_content=parsed_data.image_content,
             map_data=parsed_data.map_data,
+            room_names=room_names,
             raw_api_response=response,
         )
diff --git a/roborock/map/b01_map_parser.py b/roborock/map/b01_map_parser.py
@@ -1,15 +1,13 @@
 """Module for parsing B01/Q7 map content.
 
-B01/Q7 map uploads arrive as encrypted/compressed bytes in MQTT MAP_RESPONSE payloads.
-This module provides a small parser that:
-
-1) Decodes the raw MAP_RESPONSE payload (base64/AES/zlib layers)
-2) Parses the inflated SCMap protobuf-like payload
-3) Renders a basic map image
-
-This is intentionally minimal: it aims to provide a stable `map_content` surface
-similar to the existing v1 map content trait, without inventing additional public
-APIs.
+Observed Q7 MAP_RESPONSE payloads are:
+- base64-encoded ASCII
+- AES-ECB encrypted with the derived map key
+- PKCS7 padded
+- ASCII hex for a zlib-compressed SCMap payload
+
+This module keeps the decode path narrow and explicit to match the observed
+payload shape as closely as possible.
 """
 
 from __future__ import annotations
@@ -22,16 +20,15 @@
 from dataclasses import dataclass
 
 from Crypto.Cipher import AES
-from Crypto.Util.Padding import unpad
+from Crypto.Util.Padding import pad, unpad
 from PIL import Image
 from vacuum_map_parser_base.config.image_config import ImageConfig
-from vacuum_map_parser_base.map_data import ImageData, MapData, Room
+from vacuum_map_parser_base.map_data import ImageData, MapData
 
 from roborock.exceptions import RoborockException
 
 from .map_parser import ParsedMapData
 
-_B01_HASH = "5wwh9ikChRjASpMU8cxg7o1d2E"
 _B64_CHARS = set(b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=")
 _MAP_FILE_FORMAT = "PNG"
 
@@ -52,8 +49,9 @@ def __init__(self, config: B01MapParserConfig | None = None) -> None:
 
     def parse(self, raw_payload: bytes, *, local_key: str, serial: str, model: str) -> ParsedMapData:
         """Parse a raw MAP_RESPONSE payload and return a PNG + MapData."""
-        inflated = _decode_b01_map_payload(raw_payload, local_key=local_key, serial=serial, model=model)
-        size_x, size_y, grid, rooms = _parse_scmap_payload(inflated)
+        _ = local_key  # Current observed Q7 MAP_RESPONSE payloads do not use the local key.
+        inflated = _decode_b01_map_payload(raw_payload, serial=serial, model=model)
+        size_x, size_y, grid, room_names = _parse_scmap_payload(inflated)
 
         image = _render_occupancy_image(grid, size_x=size_x, size_y=size_y, scale=self._config.map_scale)
 
@@ -68,7 +66,8 @@ def parse(self, raw_payload: bytes, *, local_key: str, serial: str, model: str)
             data=image,
             img_transformation=lambda p: p,
         )
-        map_data.rooms = {rid: Room(0, 0, 0, 0, number=rid, name=name) for rid, name in rooms.items()} or None
+        if room_names:
+            map_data.additional_parameters["room_names"] = room_names
 
         image_bytes = io.BytesIO()
         image.save(image_bytes, format=_MAP_FILE_FORMAT)
@@ -79,112 +78,43 @@ def parse(self, raw_payload: bytes, *, local_key: str, serial: str, model: str)
         )
 
 
-def _derive_b01_iv(iv_seed: int) -> bytes:
-    random_hex = iv_seed.to_bytes(4, "big").hex().lower()
-    md5 = hashlib.md5((random_hex + _B01_HASH).encode(), usedforsecurity=False).hexdigest()
-    return md5[9:25].encode()
-
-
 def _derive_map_key(serial: str, model: str) -> bytes:
-    """Derive map decrypt key for B01/Q7 map payloads."""
+    """Derive the B01/Q7 map decrypt key from serial + model."""
     model_suffix = model.split(".")[-1]
     model_key = (model_suffix + "0" * 16)[:16].encode()
     material = f"{serial}+{model_suffix}+{serial}".encode()
-    encrypted = AES.new(model_key, AES.MODE_ECB).encrypt(_pad_pkcs7(material, AES.block_size))
+    encrypted = AES.new(model_key, AES.MODE_ECB).encrypt(pad(material, AES.block_size))
     md5 = hashlib.md5(base64.b64encode(encrypted), usedforsecurity=False).hexdigest()
     return md5[8:24].encode()
 
 
-def _pad_pkcs7(data: bytes, block_size: int) -> bytes:
-    pad_len = block_size - (len(data) % block_size)
-    return data + bytes([pad_len]) * pad_len
-
-
-def _maybe_b64(data: bytes) -> bytes | None:
-    """Decode base64 only when the input clearly looks like base64 ASCII."""
-    blob = data.strip()
-    if len(blob) < 32:
-        return None
-    if any(b not in _B64_CHARS for b in blob):
-        return None
+def _decode_base64_payload(raw_payload: bytes) -> bytes:
+    blob = raw_payload.strip()
+    if len(blob) < 32 or any(b not in _B64_CHARS for b in blob):
+        raise RoborockException("Failed to decode B01 map payload")
 
     padded = blob + b"=" * (-len(blob) % 4)
     try:
         return base64.b64decode(padded, validate=True)
-    except binascii.Error:
-        return None
+    except binascii.Error as err:
+        raise RoborockException("Failed to decode B01 map payload") from err
 
 
-def _decode_b01_map_payload(raw_payload: bytes, *, local_key: str, serial: str, model: str) -> bytes:
+def _decode_b01_map_payload(raw_payload: bytes, *, serial: str, model: str) -> bytes:
     """Decode raw B01 MAP_RESPONSE payload into inflated SCMap bytes."""
-
-    # Try progressively un-base64 layers (some payloads are double-base64).
-    layers: list[bytes] = [raw_payload]
-    l0 = _maybe_b64(raw_payload)
-    if l0 is not None:
-        layers.append(l0)
-        l1 = _maybe_b64(l0)
-        if l1 is not None:
-            layers.append(l1)
+    encrypted_payload = _decode_base64_payload(raw_payload)
+    if len(encrypted_payload) % AES.block_size != 0:
+        raise RoborockException("Unexpected encrypted B01 map payload length")
 
     map_key = _derive_map_key(serial, model)
+    decrypted_hex = AES.new(map_key, AES.MODE_ECB).decrypt(encrypted_payload)
 
-    for layer in layers:
-        candidates: list[bytes] = [layer]
-
-        # Optional B01 envelope for local-key CBC decryption.
-        if len(layer) > 19 and layer[:3] == b"B01":
-            iv_seed = int.from_bytes(layer[7:11], "big")
-            payload_len = int.from_bytes(layer[17:19], "big")
-            if payload_len and (19 + payload_len) <= len(layer):
-                encrypted = layer[19 : 19 + payload_len]
-                try:
-                    decrypted = AES.new(local_key.encode(), AES.MODE_CBC, _derive_b01_iv(iv_seed)).decrypt(encrypted)
-                    candidates.append(unpad(decrypted, AES.block_size))
-                except Exception:
-                    pass
-
-        # Optional map-key ECB layer.
-        for candidate in list(candidates):
-            if len(candidate) % AES.block_size != 0:
-                continue
-            try:
-                decrypted = AES.new(map_key, AES.MODE_ECB).decrypt(candidate)
-            except Exception:
-                continue
-
-            candidates.append(decrypted)
-            try:
-                candidates.append(unpad(decrypted, AES.block_size))
-            except Exception:
-                pass
-
-        for candidate in candidates:
-            for variant in _candidate_variants(candidate):
-                try:
-                    return zlib.decompress(variant)
-                except zlib.error:
-                    continue
-
-    raise RoborockException("Failed to decode B01 map payload")
-
-
-def _candidate_variants(candidate: bytes) -> list[bytes]:
-    variants = [candidate]
-
-    # Some payloads are ASCII hex for a compressed buffer.
     try:
-        text = candidate.strip().decode("ascii")
-    except UnicodeDecodeError:
-        text = ""
-
-    if text and len(text) >= 32 and len(text) % 2 == 0 and all(c in "0123456789abcdefABCDEF" for c in text):
-        try:
-            variants.append(bytes.fromhex(text))
-        except ValueError:
-            pass
-
-    return variants
+        compressed_hex = unpad(decrypted_hex, AES.block_size).decode("ascii")
+        compressed_payload = bytes.fromhex(compressed_hex)
+        return zlib.decompress(compressed_payload)
+    except (ValueError, UnicodeDecodeError, zlib.error) as err:
+        raise RoborockException("Failed to decode B01 map payload") from err
 
 
 def _read_varint(buf: bytes, idx: int) -> tuple[int, int]:
@@ -265,7 +195,7 @@ def _parse_scmap_payload(payload: bytes) -> tuple[int, int, bytes, dict[int, str
     size_x = 0
     size_y = 0
     grid = b""
-    rooms: dict[int, str] = {}
+    room_names: dict[int, str] = {}
 
     idx = 0
     while idx < len(payload):
@@ -308,7 +238,7 @@ def _parse_scmap_payload(payload: bytes) -> tuple[int, int, bytes, dict[int, str
         elif field_no == 12:  # roomDataInfo (repeated)
             room_id, room_name = _parse_room_data_info(value)
             if room_id is not None:
-                rooms[room_id] = room_name or f"Room {room_id}"
+                room_names[room_id] = room_name or f"Room {room_id}"
 
     if not size_x or not size_y or not grid:
         raise RoborockException("Failed to parse B01 map header/grid")
@@ -317,21 +247,20 @@ def _parse_scmap_payload(payload: bytes) -> tuple[int, int, bytes, dict[int, str
     if len(grid) < expected_len:
         raise RoborockException("B01 map data shorter than expected dimensions")
 
-    return size_x, size_y, grid[:expected_len], rooms
+    return size_x, size_y, grid[:expected_len], room_names
 
 
 def _render_occupancy_image(grid: bytes, *, size_x: int, size_y: int, scale: int) -> Image.Image:
     """Render the B01 occupancy grid into a simple image."""
 
-    # The raw grid contains occupancy-like bytes. For now we render:
-    # - 0: outside/unknown (black)
-    # - 127: wall/obstacle (grey)
-    # - 128+: floor/free (white)
+    # The observed occupancy grid contains only:
+    # - 0: outside/unknown
+    # - 127: wall/obstacle
+    # - 128: floor/free
     table = bytearray(range(256))
     table[0] = 0
     table[127] = 180
-    for i in range(128, 256):
-        table[i] = 255
+    table[128] = 255
 
     mapped = bytes(table[b] for b in grid)
     img = Image.frombytes("L", (size_x, size_y), mapped)
diff --git a/tests/devices/traits/b01/q7/test_map_content.py b/tests/devices/traits/b01/q7/test_map_content.py
@@ -27,6 +27,7 @@ async def test_q7_map_content_refresh_populates_cached_values(
     )
 
     dummy_map_data = MapData()
+    dummy_map_data.additional_parameters["room_names"] = {17: "room8"}
     with patch(
         "roborock.devices.traits.b01.q7.map_content.B01MapParser.parse",
         return_value=type("X", (), {"image_content": b"pngbytes", "map_data": dummy_map_data})(),
@@ -35,6 +36,7 @@ async def test_q7_map_content_refresh_populates_cached_values(
 
     assert q7_api.map_content.image_content == b"pngbytes"
     assert q7_api.map_content.map_data is dummy_map_data
+    assert q7_api.map_content.room_names == {17: "room8"}
     assert q7_api.map_content.raw_api_response == b"raw-map-payload"
 
     parse.assert_called_once()
@@ -44,3 +46,27 @@ def test_q7_map_content_parse_errors_cleanly(q7_api: Q7PropertiesApi):
     with patch("roborock.devices.traits.b01.q7.map_content.B01MapParser.parse", side_effect=ValueError("boom")):
         with pytest.raises(RoborockException, match="Failed to parse B01 map data"):
             q7_api.map_content.parse_map_content(b"raw")
+
+
+def test_q7_map_content_missing_metadata_fails_lazily(fake_channel: FakeChannel):
+    from roborock.data import HomeDataDevice, HomeDataProduct, RoborockCategory
+
+    q7_api = Q7PropertiesApi(
+        fake_channel,
+        device=HomeDataDevice(
+            duid="abc123",
+            name="Q7",
+            local_key="key123key123key1",
+            product_id="product-id-q7",
+            sn=None,
+        ),
+        product=HomeDataProduct(
+            id="product-id-q7",
+            name="Roborock Q7",
+            model="roborock.vacuum.sc05",
+            category=RoborockCategory.VACUUM,
+        ),
+    )
+
+    with pytest.raises(RoborockException, match="requires device serial number and model metadata"):
+        q7_api.map_content.parse_map_content(b"raw")
diff --git a/tests/map/test_b01_map_parser.py b/tests/map/test_b01_map_parser.py
@@ -34,19 +34,28 @@ def test_b01_map_parser_decodes_and_renders_fixture() -> None:
     compressed = zlib.compress(inflated)
     map_key = _derive_map_key(serial, model)
     encrypted = AES.new(map_key, AES.MODE_ECB).encrypt(pad(compressed.hex().encode(), AES.block_size))
-    payload = base64.b64encode(base64.b64encode(encrypted))
+    payload = base64.b64encode(encrypted)
 
     parser = B01MapParser()
     parsed = parser.parse(payload, local_key=local_key, serial=serial, model=model)
 
     assert parsed.image_content is not None
     assert parsed.image_content.startswith(b"\x89PNG\r\n\x1a\n")
     assert parsed.map_data is not None
-    assert parsed.map_data.rooms is not None
 
     # The fixture includes 10 rooms with names room1..room10.
-    assert set(parsed.map_data.rooms.keys()) == {10, 11, 12, 13, 14, 15, 16, 17, 18, 19}
-    assert parsed.map_data.rooms[10].name == "room1"
+    assert parsed.map_data.additional_parameters["room_names"] == {
+        10: "room1",
+        11: "room2",
+        12: "room3",
+        13: "room4",
+        14: "room5",
+        15: "room6",
+        16: "room7",
+        17: "room8",
+        18: "room9",
+        19: "room10",
+    }
 
     # Image should be scaled by default.
     img = Image.open(io.BytesIO(parsed.image_content))
