refactor: restore declarative q7 scmap fields

Author: arduano

roborock/map/b01_map_parser.py

@@ -6,9 +6,12 @@
 - PKCS7 padded
 - ASCII hex for a zlib-compressed SCMap payload
 
-The inner SCMap blob uses protobuf wire encoding. We keep the runtime parser
-minimal and declarative: only the fields needed to render the occupancy image
-and expose room names are described here.
+The inner SCMap blob is a protobuf-wire message. The declarative field layout
+below is taken from the reverse-engineered Q7 mobile app bundle (RR_API), so
+these fields are grounded in observed app schema/post-processing code rather
+than guessed parser branches. We still avoid shipping generated protobuf
+classes from a reverse-engineered schema, which keeps the runtime parser
+explicit without overstating certainty about the full message definition.
 """
 
 from __future__ import annotations
@@ -17,9 +20,10 @@
 import binascii
 import hashlib
 import io
+import struct
 import zlib
 from collections.abc import Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from Crypto.Cipher import AES
 from Crypto.Util.Padding import pad, unpad
@@ -49,6 +53,70 @@ class _ProtoField:
     parser: Callable[[object], object] | None = None
 
 
+@dataclass(frozen=True)
+class _ScPoint:
+    x: float | None = None
+    y: float | None = None
+
+
+@dataclass(frozen=True)
+class _ScMapBoundaryInfo:
+    map_md5: str | None = None
+    v_min_x: int | None = None
+    v_max_x: int | None = None
+    v_min_y: int | None = None
+    v_max_y: int | None = None
+
+
+@dataclass(frozen=True)
+class _ScMapExtInfo:
+    task_begin_date: int | None = None
+    map_upload_date: int | None = None
+    map_valid: int | None = None
+    radian: int | None = None
+    force: int | None = None
+    clean_path: int | None = None
+    boundary_info: _ScMapBoundaryInfo | None = None
+    map_version: int | None = None
+    map_value_type: int | None = None
+
+
+@dataclass(frozen=True)
+class _ScMapHead:
+    map_head_id: int | None = None
+    size_x: int | None = None
+    size_y: int | None = None
+    min_x: float | None = None
+    min_y: float | None = None
+    max_x: float | None = None
+    max_y: float | None = None
+    resolution: float | None = None
+
+
+@dataclass(frozen=True)
+class _ScRoomData:
+    room_id: int | None = None
+    room_name: str | None = None
+    room_type_id: int | None = None
+    material_id: int | None = None
+    clean_state: int | None = None
+    room_clean: int | None = None
+    room_clean_index: int | None = None
+    room_name_post: _ScPoint | None = None
+    color_id: int | None = None
+    floor_direction: int | None = None
+    global_seq: int | None = None
+
+
+@dataclass(frozen=True)
+class _ScMapPayload:
+    map_type: int | None = None
+    map_ext_info: _ScMapExtInfo | None = None
+    map_head: _ScMapHead | None = None
+    map_data: bytes | None = None
+    room_data_info: tuple[_ScRoomData, ...] = field(default_factory=tuple)
+
+
 @dataclass
 class B01MapParserConfig:
     """Configuration for the B01/Q7 map parser."""
@@ -66,7 +134,9 @@ def __init__(self, config: B01MapParserConfig | None = None) -> None:
     def parse(self, raw_payload: bytes, *, serial: str, model: str) -> ParsedMapData:
         """Parse a raw MAP_RESPONSE payload and return a PNG + MapData."""
         inflated = _decode_b01_map_payload(raw_payload, serial=serial, model=model)
-        size_x, size_y, grid, room_names = _parse_scmap_payload(inflated)
+        scmap = _parse_scmap_payload(inflated)
+        size_x, size_y, grid = _extract_grid(scmap)
+        room_names = _extract_room_names(scmap.room_data_info)
 
         image = _render_occupancy_image(grid, size_x=size_x, size_y=size_y, scale=self._config.map_scale)
 
@@ -178,6 +248,10 @@ def _decode_utf8(value: bytes) -> str:
     return value.decode("utf-8", errors="replace")
 
 
+def _decode_float32(value: bytes) -> float:
+    return struct.unpack("<f", value)[0]
+
+
 def _decode_map_data_bytes(value: bytes) -> bytes:
     try:
         return zlib.decompress(value)
@@ -218,47 +292,145 @@ def _parse_proto_message(blob: bytes, schema: dict[int, _ProtoField], *, context
     return parsed
 
 
-def _parse_map_head_info(blob: bytes) -> tuple[int | None, int | None]:
-    parsed = _parse_proto_message(blob, _MAP_HEAD_INFO_SCHEMA, context="B01 MapHeadInfo")
-    return parsed.get("size_x"), parsed.get("size_y")
+def _parse_sc_point(blob: bytes) -> _ScPoint:
+    parsed = _parse_proto_message(blob, _DEVICE_POINT_INFO_SCHEMA, context="B01 DevicePointInfo")
+    return _ScPoint(x=parsed.get("x"), y=parsed.get("y"))
+
+
+def _parse_sc_map_boundary_info(blob: bytes) -> _ScMapBoundaryInfo:
+    parsed = _parse_proto_message(blob, _MAP_BOUNDARY_INFO_SCHEMA, context="B01 MapBoundaryInfo")
+    return _ScMapBoundaryInfo(
+        map_md5=parsed.get("map_md5"),
+        v_min_x=parsed.get("v_min_x"),
+        v_max_x=parsed.get("v_max_x"),
+        v_min_y=parsed.get("v_min_y"),
+        v_max_y=parsed.get("v_max_y"),
+    )
 
 
-def _parse_map_data_info(blob: bytes) -> bytes:
+def _parse_sc_map_ext_info(blob: bytes) -> _ScMapExtInfo:
+    parsed = _parse_proto_message(blob, _MAP_EXT_INFO_SCHEMA, context="B01 MapExtInfo")
+    return _ScMapExtInfo(
+        task_begin_date=parsed.get("task_begin_date"),
+        map_upload_date=parsed.get("map_upload_date"),
+        map_valid=parsed.get("map_valid"),
+        radian=parsed.get("radian"),
+        force=parsed.get("force"),
+        clean_path=parsed.get("clean_path"),
+        boundary_info=parsed.get("boundary_info"),
+        map_version=parsed.get("map_version"),
+        map_value_type=parsed.get("map_value_type"),
+    )
+
+
+def _parse_sc_map_head(blob: bytes) -> _ScMapHead:
+    parsed = _parse_proto_message(blob, _MAP_HEAD_INFO_SCHEMA, context="B01 MapHeadInfo")
+    return _ScMapHead(
+        map_head_id=parsed.get("map_head_id"),
+        size_x=parsed.get("size_x"),
+        size_y=parsed.get("size_y"),
+        min_x=parsed.get("min_x"),
+        min_y=parsed.get("min_y"),
+        max_x=parsed.get("max_x"),
+        max_y=parsed.get("max_y"),
+        resolution=parsed.get("resolution"),
+    )
+
+
+def _parse_sc_map_data_info(blob: bytes) -> bytes:
     parsed = _parse_proto_message(blob, _MAP_DATA_INFO_SCHEMA, context="B01 MapDataInfo")
     if (map_data := parsed.get("map_data")) is None:
         raise RoborockException("B01 map payload missing mapData")
     return map_data
 
 
-def _parse_room_data_info(blob: bytes) -> tuple[int | None, str | None]:
+def _parse_sc_room_data(blob: bytes) -> _ScRoomData:
     parsed = _parse_proto_message(blob, _ROOM_DATA_INFO_SCHEMA, context="B01 RoomDataInfo")
-    return parsed.get("room_id"), parsed.get("room_name")
+    return _ScRoomData(
+        room_id=parsed.get("room_id"),
+        room_name=parsed.get("room_name"),
+        room_type_id=parsed.get("room_type_id"),
+        material_id=parsed.get("material_id"),
+        clean_state=parsed.get("clean_state"),
+        room_clean=parsed.get("room_clean"),
+        room_clean_index=parsed.get("room_clean_index"),
+        room_name_post=parsed.get("room_name_post"),
+        color_id=parsed.get("color_id"),
+        floor_direction=parsed.get("floor_direction"),
+        global_seq=parsed.get("global_seq"),
+    )
+
+
+def _parse_scmap_payload(payload: bytes) -> _ScMapPayload:
+    """Parse inflated SCMap bytes using the reverse-engineered app field layout."""
+    parsed = _parse_proto_message(payload, _ROBOT_MAP_SCHEMA, context="B01 SCMap")
+    return _ScMapPayload(
+        map_type=parsed.get("map_type"),
+        map_ext_info=parsed.get("map_ext_info"),
+        map_head=parsed.get("map_head"),
+        map_data=parsed.get("map_data"),
+        room_data_info=tuple(parsed.get("room_data_info", [])),
+    )
 
 
-def _parse_scmap_payload(payload: bytes) -> tuple[int, int, bytes, dict[int, str]]:
-    """Parse inflated SCMap bytes using the reverse-engineered field layout we need."""
-    parsed = _parse_proto_message(payload, _ROBOT_MAP_SCHEMA, context="B01 SCMap")
+def _extract_grid(scmap: _ScMapPayload) -> tuple[int, int, bytes]:
+    if scmap.map_head is None or scmap.map_data is None:
+        raise RoborockException("Failed to parse B01 map header/grid")
 
-    size_x, size_y = parsed.get("map_head", (None, None))
-    grid = parsed.get("map_data")
-    if not size_x or not size_y or not grid:
+    size_x = scmap.map_head.size_x or 0
+    size_y = scmap.map_head.size_y or 0
+    if not size_x or not size_y or not scmap.map_data:
         raise RoborockException("Failed to parse B01 map header/grid")
 
     expected_len = size_x * size_y
-    if len(grid) < expected_len:
+    if len(scmap.map_data) < expected_len:
         raise RoborockException("B01 map data shorter than expected dimensions")
 
+    return size_x, size_y, scmap.map_data[:expected_len]
+
+
+def _extract_room_names(rooms: tuple[_ScRoomData, ...]) -> dict[int, str]:
     room_names: dict[int, str] = {}
-    for room_id, room_name in parsed.get("room_data_info", []):
-        if room_id is not None:
-            room_names[room_id] = room_name or f"Room {room_id}"
+    for room in rooms:
+        if room.room_id is not None:
+            room_names[room.room_id] = room.room_name or f"Room {room.room_id}"
+    return room_names
 
-    return size_x, size_y, grid[:expected_len], room_names
 
+_DEVICE_POINT_INFO_SCHEMA = {
+    1: _ProtoField("x", _WIRE_FIXED32, parser=_decode_float32),
+    2: _ProtoField("y", _WIRE_FIXED32, parser=_decode_float32),
+}
+
+_MAP_BOUNDARY_INFO_SCHEMA = {
+    1: _ProtoField("map_md5", _WIRE_LEN, parser=_decode_utf8),
+    2: _ProtoField("v_min_x", _WIRE_VARINT, parser=_decode_uint32),
+    3: _ProtoField("v_max_x", _WIRE_VARINT, parser=_decode_uint32),
+    4: _ProtoField("v_min_y", _WIRE_VARINT, parser=_decode_uint32),
+    5: _ProtoField("v_max_y", _WIRE_VARINT, parser=_decode_uint32),
+}
+
+_MAP_EXT_INFO_SCHEMA = {
+    1: _ProtoField("task_begin_date", _WIRE_VARINT, parser=_decode_uint32),
+    2: _ProtoField("map_upload_date", _WIRE_VARINT, parser=_decode_uint32),
+    3: _ProtoField("map_valid", _WIRE_VARINT, parser=_decode_uint32),
+    4: _ProtoField("radian", _WIRE_VARINT, parser=_decode_uint32),
+    5: _ProtoField("force", _WIRE_VARINT, parser=_decode_uint32),
+    6: _ProtoField("clean_path", _WIRE_VARINT, parser=_decode_uint32),
+    7: _ProtoField("boundary_info", _WIRE_LEN, parser=_parse_sc_map_boundary_info),
+    8: _ProtoField("map_version", _WIRE_VARINT, parser=_decode_uint32),
+    9: _ProtoField("map_value_type", _WIRE_VARINT, parser=_decode_uint32),
+}
 
 _MAP_HEAD_INFO_SCHEMA = {
+    1: _ProtoField("map_head_id", _WIRE_VARINT, parser=_decode_uint32),
     2: _ProtoField("size_x", _WIRE_VARINT, parser=_decode_uint32),
     3: _ProtoField("size_y", _WIRE_VARINT, parser=_decode_uint32),
+    4: _ProtoField("min_x", _WIRE_FIXED32, parser=_decode_float32),
+    5: _ProtoField("min_y", _WIRE_FIXED32, parser=_decode_float32),
+    6: _ProtoField("max_x", _WIRE_FIXED32, parser=_decode_float32),
+    7: _ProtoField("max_y", _WIRE_FIXED32, parser=_decode_float32),
+    8: _ProtoField("resolution", _WIRE_FIXED32, parser=_decode_float32),
 }
 
 _MAP_DATA_INFO_SCHEMA = {
@@ -268,12 +440,23 @@ def _parse_scmap_payload(payload: bytes) -> tuple[int, int, bytes, dict[int, str
 _ROOM_DATA_INFO_SCHEMA = {
     1: _ProtoField("room_id", _WIRE_VARINT, parser=_decode_uint32),
     2: _ProtoField("room_name", _WIRE_LEN, parser=_decode_utf8),
+    3: _ProtoField("room_type_id", _WIRE_VARINT, parser=_decode_uint32),
+    4: _ProtoField("material_id", _WIRE_VARINT, parser=_decode_uint32),
+    5: _ProtoField("clean_state", _WIRE_VARINT, parser=_decode_uint32),
+    6: _ProtoField("room_clean", _WIRE_VARINT, parser=_decode_uint32),
+    7: _ProtoField("room_clean_index", _WIRE_VARINT, parser=_decode_uint32),
+    8: _ProtoField("room_name_post", _WIRE_LEN, parser=_parse_sc_point),
+    10: _ProtoField("color_id", _WIRE_VARINT, parser=_decode_uint32),
+    11: _ProtoField("floor_direction", _WIRE_VARINT, parser=_decode_uint32),
+    12: _ProtoField("global_seq", _WIRE_VARINT, parser=_decode_uint32),
 }
 
 _ROBOT_MAP_SCHEMA = {
-    3: _ProtoField("map_head", _WIRE_LEN, parser=_parse_map_head_info),
-    4: _ProtoField("map_data", _WIRE_LEN, parser=_parse_map_data_info),
-    12: _ProtoField("room_data_info", _WIRE_LEN, repeated=True, parser=_parse_room_data_info),
+    1: _ProtoField("map_type", _WIRE_VARINT, parser=_decode_uint32),
+    2: _ProtoField("map_ext_info", _WIRE_LEN, parser=_parse_sc_map_ext_info),
+    3: _ProtoField("map_head", _WIRE_LEN, parser=_parse_sc_map_head),
+    4: _ProtoField("map_data", _WIRE_LEN, parser=_parse_sc_map_data_info),
+    12: _ProtoField("room_data_info", _WIRE_LEN, repeated=True, parser=_parse_sc_room_data),
 }