Skip to content

Commit bedf379

Browse files
arduanoarduano
committed
feat: define checked-in proto for q7 scmap
1 parent 2d73fd1 commit bedf379

File tree

5 files changed

+181
-215
lines changed

5 files changed

+181
-215
lines changed

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ line-length = 120
102102

103103
[tool.ruff.lint.per-file-ignores]
104104
"*/__init__.py" = ["F401"]
105+
"roborock/map/proto/*_pb2.py" = ["E501", "I001", "UP009"]
105106

106107
[tool.pytest.ini_options]
107108
asyncio_mode = "auto"

‎roborock/map/b01_map_parser.py‎

Lines changed: 59 additions & 215 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@
66
- PKCS7 padded
77
- ASCII hex for a zlib-compressed SCMap payload
88
9-
The inner SCMap blob is parsed with the official protobuf runtime using a small
10-
runtime descriptor for the message fields this parser needs.
9+
The inner SCMap blob is parsed with protobuf messages generated from
10+
`roborock/map/proto/b01_scmap.proto`.
1111
"""
1212

1313
from __future__ import annotations
@@ -18,25 +18,21 @@
1818
import io
1919
import zlib
2020
from dataclasses import dataclass, field
21-
from typing import Any
2221

2322
from Crypto.Cipher import AES
2423
from Crypto.Util.Padding import pad, unpad
25-
from google.protobuf import descriptor_pool
26-
from google.protobuf.descriptor_pb2 import DescriptorProto, FieldDescriptorProto, FileDescriptorProto
27-
from google.protobuf.message import DecodeError, Message
28-
from google.protobuf.message_factory import GetMessageClass
24+
from google.protobuf.message import DecodeError
2925
from PIL import Image
3026
from vacuum_map_parser_base.config.image_config import ImageConfig
3127
from vacuum_map_parser_base.map_data import ImageData, MapData
3228

3329
from roborock.exceptions import RoborockException
30+
from roborock.map.proto import b01_scmap_pb2
3431

3532
from .map_parser import ParsedMapData
3633

3734
_B64_CHARS = set(b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=")
3835
_MAP_FILE_FORMAT = "PNG"
39-
_PROTO_PACKAGE = "b01.scmap"
4036

4137

4238
@dataclass(frozen=True)
@@ -188,151 +184,11 @@ def _decode_b01_map_payload(raw_payload: bytes, *, serial: str, model: str) -> b
188184
raise RoborockException("Failed to decode B01 map payload") from err
189185

190186

191-
def _message_descriptor(name: str, fields: list[dict[str, object]]) -> DescriptorProto:
192-
descriptor = DescriptorProto(name=name)
193-
for field_def in fields:
194-
descriptor.field.add(
195-
name=field_def["name"],
196-
number=field_def["number"],
197-
label=field_def.get("label", FieldDescriptorProto.LABEL_OPTIONAL),
198-
type=field_def["type"],
199-
type_name=field_def.get("type_name"),
200-
)
201-
return descriptor
202-
203-
204-
_FILE_DESCRIPTOR = FileDescriptorProto(name="b01_scmap.proto", package=_PROTO_PACKAGE, syntax="proto2")
205-
_FILE_DESCRIPTOR.message_type.extend(
206-
[
207-
_message_descriptor(
208-
"DevicePointInfo",
209-
[
210-
{"name": "x", "number": 1, "type": FieldDescriptorProto.TYPE_FLOAT},
211-
{"name": "y", "number": 2, "type": FieldDescriptorProto.TYPE_FLOAT},
212-
],
213-
),
214-
_message_descriptor(
215-
"MapBoundaryInfo",
216-
[
217-
{"name": "mapMd5", "number": 1, "type": FieldDescriptorProto.TYPE_STRING},
218-
{"name": "vMinX", "number": 2, "type": FieldDescriptorProto.TYPE_UINT32},
219-
{"name": "vMaxX", "number": 3, "type": FieldDescriptorProto.TYPE_UINT32},
220-
{"name": "vMinY", "number": 4, "type": FieldDescriptorProto.TYPE_UINT32},
221-
{"name": "vMaxY", "number": 5, "type": FieldDescriptorProto.TYPE_UINT32},
222-
],
223-
),
224-
_message_descriptor(
225-
"MapExtInfo",
226-
[
227-
{"name": "taskBeginDate", "number": 1, "type": FieldDescriptorProto.TYPE_UINT32},
228-
{"name": "mapUploadDate", "number": 2, "type": FieldDescriptorProto.TYPE_UINT32},
229-
{"name": "mapValid", "number": 3, "type": FieldDescriptorProto.TYPE_UINT32},
230-
{"name": "radian", "number": 4, "type": FieldDescriptorProto.TYPE_UINT32},
231-
{"name": "force", "number": 5, "type": FieldDescriptorProto.TYPE_UINT32},
232-
{"name": "cleanPath", "number": 6, "type": FieldDescriptorProto.TYPE_UINT32},
233-
{
234-
"name": "boudaryInfo",
235-
"number": 7,
236-
"type": FieldDescriptorProto.TYPE_MESSAGE,
237-
"type_name": f".{_PROTO_PACKAGE}.MapBoundaryInfo",
238-
},
239-
{"name": "mapVersion", "number": 8, "type": FieldDescriptorProto.TYPE_UINT32},
240-
{"name": "mapValueType", "number": 9, "type": FieldDescriptorProto.TYPE_UINT32},
241-
],
242-
),
243-
_message_descriptor(
244-
"MapHeadInfo",
245-
[
246-
{"name": "mapHeadId", "number": 1, "type": FieldDescriptorProto.TYPE_UINT32},
247-
{"name": "sizeX", "number": 2, "type": FieldDescriptorProto.TYPE_UINT32},
248-
{"name": "sizeY", "number": 3, "type": FieldDescriptorProto.TYPE_UINT32},
249-
{"name": "minX", "number": 4, "type": FieldDescriptorProto.TYPE_FLOAT},
250-
{"name": "minY", "number": 5, "type": FieldDescriptorProto.TYPE_FLOAT},
251-
{"name": "maxX", "number": 6, "type": FieldDescriptorProto.TYPE_FLOAT},
252-
{"name": "maxY", "number": 7, "type": FieldDescriptorProto.TYPE_FLOAT},
253-
{"name": "resolution", "number": 8, "type": FieldDescriptorProto.TYPE_FLOAT},
254-
],
255-
),
256-
_message_descriptor(
257-
"MapDataInfo",
258-
[{"name": "mapData", "number": 1, "type": FieldDescriptorProto.TYPE_BYTES}],
259-
),
260-
_message_descriptor(
261-
"RoomDataInfo",
262-
[
263-
{"name": "roomId", "number": 1, "type": FieldDescriptorProto.TYPE_UINT32},
264-
{"name": "roomName", "number": 2, "type": FieldDescriptorProto.TYPE_STRING},
265-
{"name": "roomTypeId", "number": 3, "type": FieldDescriptorProto.TYPE_UINT32},
266-
{"name": "meterialId", "number": 4, "type": FieldDescriptorProto.TYPE_UINT32},
267-
{"name": "cleanState", "number": 5, "type": FieldDescriptorProto.TYPE_UINT32},
268-
{"name": "roomClean", "number": 6, "type": FieldDescriptorProto.TYPE_UINT32},
269-
{"name": "roomCleanIndex", "number": 7, "type": FieldDescriptorProto.TYPE_UINT32},
270-
{
271-
"name": "roomNamePost",
272-
"number": 8,
273-
"type": FieldDescriptorProto.TYPE_MESSAGE,
274-
"type_name": f".{_PROTO_PACKAGE}.DevicePointInfo",
275-
},
276-
{"name": "colorId", "number": 10, "type": FieldDescriptorProto.TYPE_UINT32},
277-
{"name": "floor_direction", "number": 11, "type": FieldDescriptorProto.TYPE_UINT32},
278-
{"name": "global_seq", "number": 12, "type": FieldDescriptorProto.TYPE_UINT32},
279-
],
280-
),
281-
_message_descriptor(
282-
"RobotMap",
283-
[
284-
{"name": "mapType", "number": 1, "type": FieldDescriptorProto.TYPE_UINT32},
285-
{
286-
"name": "mapExtInfo",
287-
"number": 2,
288-
"type": FieldDescriptorProto.TYPE_MESSAGE,
289-
"type_name": f".{_PROTO_PACKAGE}.MapExtInfo",
290-
},
291-
{
292-
"name": "mapHead",
293-
"number": 3,
294-
"type": FieldDescriptorProto.TYPE_MESSAGE,
295-
"type_name": f".{_PROTO_PACKAGE}.MapHeadInfo",
296-
},
297-
{
298-
"name": "mapData",
299-
"number": 4,
300-
"type": FieldDescriptorProto.TYPE_MESSAGE,
301-
"type_name": f".{_PROTO_PACKAGE}.MapDataInfo",
302-
},
303-
{
304-
"name": "roomDataInfo",
305-
"number": 12,
306-
"label": FieldDescriptorProto.LABEL_REPEATED,
307-
"type": FieldDescriptorProto.TYPE_MESSAGE,
308-
"type_name": f".{_PROTO_PACKAGE}.RoomDataInfo",
309-
},
310-
],
311-
),
312-
]
313-
)
314-
315-
_SC_MAP_FILE_DESCRIPTOR = descriptor_pool.Default().AddSerializedFile(_FILE_DESCRIPTOR.SerializeToString())
316-
_DEVICE_POINT_INFO = GetMessageClass(_SC_MAP_FILE_DESCRIPTOR.message_types_by_name["DevicePointInfo"])
317-
_MAP_BOUNDARY_INFO = GetMessageClass(_SC_MAP_FILE_DESCRIPTOR.message_types_by_name["MapBoundaryInfo"])
318-
_MAP_EXT_INFO = GetMessageClass(_SC_MAP_FILE_DESCRIPTOR.message_types_by_name["MapExtInfo"])
319-
_MAP_HEAD_INFO = GetMessageClass(_SC_MAP_FILE_DESCRIPTOR.message_types_by_name["MapHeadInfo"])
320-
_MAP_DATA_INFO = GetMessageClass(_SC_MAP_FILE_DESCRIPTOR.message_types_by_name["MapDataInfo"])
321-
_ROOM_DATA_INFO = GetMessageClass(_SC_MAP_FILE_DESCRIPTOR.message_types_by_name["RoomDataInfo"])
322-
_ROBOT_MAP = GetMessageClass(_SC_MAP_FILE_DESCRIPTOR.message_types_by_name["RobotMap"])
323-
324-
325-
def _has_field(message: Any, field_name: str) -> bool:
326-
return message.HasField(field_name)
327-
328-
329-
def _parse_proto(blob: bytes, message_class: type[Message], *, context: str) -> Any:
330-
message = message_class()
187+
def _parse_proto(blob: bytes, message: object, *, context: str) -> None:
331188
try:
332189
message.ParseFromString(blob)
333190
except DecodeError as err:
334191
raise RoborockException(f"Failed to parse {context}") from err
335-
return message
336192

337193

338194
def _decode_map_data_bytes(value: bytes) -> bytes:
@@ -342,95 +198,83 @@ def _decode_map_data_bytes(value: bytes) -> bytes:
342198
return value
343199

344200

345-
def _parse_sc_point(blob: bytes) -> _ScPoint:
346-
parsed = _parse_proto(blob, _DEVICE_POINT_INFO, context="B01 DevicePointInfo")
201+
def _parse_sc_point(parsed: b01_scmap_pb2.DevicePointInfo) -> _ScPoint:
347202
return _ScPoint(
348-
x=parsed.x if _has_field(parsed, "x") else None,
349-
y=parsed.y if _has_field(parsed, "y") else None,
203+
x=parsed.x if parsed.HasField("x") else None,
204+
y=parsed.y if parsed.HasField("y") else None,
350205
)
351206

352207

353-
def _parse_sc_map_boundary_info(blob: bytes) -> _ScMapBoundaryInfo:
354-
parsed = _parse_proto(blob, _MAP_BOUNDARY_INFO, context="B01 MapBoundaryInfo")
208+
def _parse_sc_map_boundary_info(parsed: b01_scmap_pb2.MapBoundaryInfo) -> _ScMapBoundaryInfo:
355209
return _ScMapBoundaryInfo(
356-
map_md5=parsed.mapMd5 if _has_field(parsed, "mapMd5") else None,
357-
v_min_x=parsed.vMinX if _has_field(parsed, "vMinX") else None,
358-
v_max_x=parsed.vMaxX if _has_field(parsed, "vMaxX") else None,
359-
v_min_y=parsed.vMinY if _has_field(parsed, "vMinY") else None,
360-
v_max_y=parsed.vMaxY if _has_field(parsed, "vMaxY") else None,
210+
map_md5=parsed.mapMd5 if parsed.HasField("mapMd5") else None,
211+
v_min_x=parsed.vMinX if parsed.HasField("vMinX") else None,
212+
v_max_x=parsed.vMaxX if parsed.HasField("vMaxX") else None,
213+
v_min_y=parsed.vMinY if parsed.HasField("vMinY") else None,
214+
v_max_y=parsed.vMaxY if parsed.HasField("vMaxY") else None,
361215
)
362216

363217

364-
def _parse_sc_map_ext_info(blob: bytes) -> _ScMapExtInfo:
365-
parsed = _parse_proto(blob, _MAP_EXT_INFO, context="B01 MapExtInfo")
218+
def _parse_sc_map_ext_info(parsed: b01_scmap_pb2.MapExtInfo) -> _ScMapExtInfo:
366219
return _ScMapExtInfo(
367-
task_begin_date=parsed.taskBeginDate if _has_field(parsed, "taskBeginDate") else None,
368-
map_upload_date=parsed.mapUploadDate if _has_field(parsed, "mapUploadDate") else None,
369-
map_valid=parsed.mapValid if _has_field(parsed, "mapValid") else None,
370-
radian=parsed.radian if _has_field(parsed, "radian") else None,
371-
force=parsed.force if _has_field(parsed, "force") else None,
372-
clean_path=parsed.cleanPath if _has_field(parsed, "cleanPath") else None,
373-
boundary_info=(
374-
_parse_sc_map_boundary_info(parsed.boudaryInfo.SerializeToString())
375-
if _has_field(parsed, "boudaryInfo")
376-
else None
377-
),
378-
map_version=parsed.mapVersion if _has_field(parsed, "mapVersion") else None,
379-
map_value_type=parsed.mapValueType if _has_field(parsed, "mapValueType") else None,
220+
task_begin_date=parsed.taskBeginDate if parsed.HasField("taskBeginDate") else None,
221+
map_upload_date=parsed.mapUploadDate if parsed.HasField("mapUploadDate") else None,
222+
map_valid=parsed.mapValid if parsed.HasField("mapValid") else None,
223+
radian=parsed.radian if parsed.HasField("radian") else None,
224+
force=parsed.force if parsed.HasField("force") else None,
225+
clean_path=parsed.cleanPath if parsed.HasField("cleanPath") else None,
226+
boundary_info=_parse_sc_map_boundary_info(parsed.boudaryInfo) if parsed.HasField("boudaryInfo") else None,
227+
map_version=parsed.mapVersion if parsed.HasField("mapVersion") else None,
228+
map_value_type=parsed.mapValueType if parsed.HasField("mapValueType") else None,
380229
)
381230

382231

383-
def _parse_sc_map_head(blob: bytes) -> _ScMapHead:
384-
parsed = _parse_proto(blob, _MAP_HEAD_INFO, context="B01 MapHeadInfo")
232+
def _parse_sc_map_head(parsed: b01_scmap_pb2.MapHeadInfo) -> _ScMapHead:
385233
return _ScMapHead(
386-
map_head_id=parsed.mapHeadId if _has_field(parsed, "mapHeadId") else None,
387-
size_x=parsed.sizeX if _has_field(parsed, "sizeX") else None,
388-
size_y=parsed.sizeY if _has_field(parsed, "sizeY") else None,
389-
min_x=parsed.minX if _has_field(parsed, "minX") else None,
390-
min_y=parsed.minY if _has_field(parsed, "minY") else None,
391-
max_x=parsed.maxX if _has_field(parsed, "maxX") else None,
392-
max_y=parsed.maxY if _has_field(parsed, "maxY") else None,
393-
resolution=parsed.resolution if _has_field(parsed, "resolution") else None,
234+
map_head_id=parsed.mapHeadId if parsed.HasField("mapHeadId") else None,
235+
size_x=parsed.sizeX if parsed.HasField("sizeX") else None,
236+
size_y=parsed.sizeY if parsed.HasField("sizeY") else None,
237+
min_x=parsed.minX if parsed.HasField("minX") else None,
238+
min_y=parsed.minY if parsed.HasField("minY") else None,
239+
max_x=parsed.maxX if parsed.HasField("maxX") else None,
240+
max_y=parsed.maxY if parsed.HasField("maxY") else None,
241+
resolution=parsed.resolution if parsed.HasField("resolution") else None,
394242
)
395243

396244

397-
def _parse_sc_map_data_info(blob: bytes) -> bytes:
398-
parsed = _parse_proto(blob, _MAP_DATA_INFO, context="B01 MapDataInfo")
399-
if not _has_field(parsed, "mapData"):
400-
raise RoborockException("B01 map payload missing mapData")
401-
return _decode_map_data_bytes(parsed.mapData)
402-
403-
404-
def _parse_sc_room_data(blob: bytes) -> _ScRoomData:
405-
parsed = _parse_proto(blob, _ROOM_DATA_INFO, context="B01 RoomDataInfo")
245+
def _parse_sc_room_data(parsed: b01_scmap_pb2.RoomDataInfo) -> _ScRoomData:
406246
return _ScRoomData(
407-
room_id=parsed.roomId if _has_field(parsed, "roomId") else None,
408-
room_name=parsed.roomName if _has_field(parsed, "roomName") else None,
409-
room_type_id=parsed.roomTypeId if _has_field(parsed, "roomTypeId") else None,
410-
material_id=parsed.meterialId if _has_field(parsed, "meterialId") else None,
411-
clean_state=parsed.cleanState if _has_field(parsed, "cleanState") else None,
412-
room_clean=parsed.roomClean if _has_field(parsed, "roomClean") else None,
413-
room_clean_index=parsed.roomCleanIndex if _has_field(parsed, "roomCleanIndex") else None,
414-
room_name_post=(
415-
_parse_sc_point(parsed.roomNamePost.SerializeToString()) if _has_field(parsed, "roomNamePost") else None
416-
),
417-
color_id=parsed.colorId if _has_field(parsed, "colorId") else None,
418-
floor_direction=parsed.floor_direction if _has_field(parsed, "floor_direction") else None,
419-
global_seq=parsed.global_seq if _has_field(parsed, "global_seq") else None,
247+
room_id=parsed.roomId if parsed.HasField("roomId") else None,
248+
room_name=parsed.roomName if parsed.HasField("roomName") else None,
249+
room_type_id=parsed.roomTypeId if parsed.HasField("roomTypeId") else None,
250+
material_id=parsed.meterialId if parsed.HasField("meterialId") else None,
251+
clean_state=parsed.cleanState if parsed.HasField("cleanState") else None,
252+
room_clean=parsed.roomClean if parsed.HasField("roomClean") else None,
253+
room_clean_index=parsed.roomCleanIndex if parsed.HasField("roomCleanIndex") else None,
254+
room_name_post=_parse_sc_point(parsed.roomNamePost) if parsed.HasField("roomNamePost") else None,
255+
color_id=parsed.colorId if parsed.HasField("colorId") else None,
256+
floor_direction=parsed.floor_direction if parsed.HasField("floor_direction") else None,
257+
global_seq=parsed.global_seq if parsed.HasField("global_seq") else None,
420258
)
421259

422260

423261
def _parse_scmap_payload(payload: bytes) -> _ScMapPayload:
424262
"""Parse inflated SCMap bytes into typed map metadata."""
425-
parsed = _parse_proto(payload, _ROBOT_MAP, context="B01 SCMap")
263+
parsed = b01_scmap_pb2.RobotMap()
264+
_parse_proto(payload, parsed, context="B01 SCMap")
265+
266+
map_data = None
267+
if parsed.HasField("mapData"):
268+
if not parsed.mapData.HasField("mapData"):
269+
raise RoborockException("B01 map payload missing mapData")
270+
map_data = _decode_map_data_bytes(parsed.mapData.mapData)
271+
426272
return _ScMapPayload(
427-
map_type=parsed.mapType if _has_field(parsed, "mapType") else None,
428-
map_ext_info=(
429-
_parse_sc_map_ext_info(parsed.mapExtInfo.SerializeToString()) if _has_field(parsed, "mapExtInfo") else None
430-
),
431-
map_head=_parse_sc_map_head(parsed.mapHead.SerializeToString()) if _has_field(parsed, "mapHead") else None,
432-
map_data=_parse_sc_map_data_info(parsed.mapData.SerializeToString()) if _has_field(parsed, "mapData") else None,
433-
room_data_info=tuple(_parse_sc_room_data(room.SerializeToString()) for room in parsed.roomDataInfo),
273+
map_type=parsed.mapType if parsed.HasField("mapType") else None,
274+
map_ext_info=_parse_sc_map_ext_info(parsed.mapExtInfo) if parsed.HasField("mapExtInfo") else None,
275+
map_head=_parse_sc_map_head(parsed.mapHead) if parsed.HasField("mapHead") else None,
276+
map_data=map_data,
277+
room_data_info=tuple(_parse_sc_room_data(room) for room in parsed.roomDataInfo),
434278
)
435279

436280

‎roborock/map/proto/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
"""Generated protobuf modules for Roborock map payloads."""

0 commit comments

Comments
 (0)