🦎 q7: tighten B01 map payload decode + use gz fixture

Author: arduano

roborock/protocols/b01_map_protocol.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import base64
+import binascii
 import hashlib
 import zlib
 
@@ -36,25 +37,42 @@ def derive_map_key(serial: str, model: str) -> bytes:
     return md5[8:24].encode()
 
 
+_B64_CHARS = set(b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/=")
+
+
 def _maybe_b64(data: bytes) -> bytes | None:
+    """Decode base64 *only* when the input clearly looks like base64 ASCII.
+
+    MAP_RESPONSE payloads sometimes arrive as base64-encoded bytes (and in some
+    cases double-base64). Avoid blindly attempting b64decode on arbitrary binary
+    data; that can yield garbage-but-decodable output and hide the real payload.
+    """
+
+    blob = data.strip()
+    if len(blob) < 32:
+        return None
+    if any(b not in _B64_CHARS for b in blob):
+        return None
+
+    # Base64 strings may omit padding.
+    padded = blob + b"=" * (-len(blob) % 4)
     try:
-        return base64.b64decode(data, validate=False)
-    except Exception:
+        return base64.b64decode(padded, validate=True)
+    except binascii.Error:
         return None
 
 
 def decode_b01_map_payload(raw_payload: bytes, *, local_key: str, serial: str, model: str) -> bytes:
     """Decode raw B01 MAP_RESPONSE payload into inflated SCMap protobuf bytes."""
 
-    layers: list[bytes] = []
+    # Try the raw payload first, then progressively un-base64 layers if present.
+    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)
-    else:
-        layers.append(raw_payload)
 
     map_key = derive_map_key(serial, model)
     for layer in layers:
@@ -64,12 +82,13 @@ def decode_b01_map_payload(raw_payload: bytes, *, local_key: str, serial: str, m
         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")
-            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, 16))
-            except Exception:
-                pass
+            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, 16))
+                except Exception:
+                    pass
 
         # Optional map-key ECB layer seen in Q7 payloads.
         for candidate in list(candidates):
@@ -84,11 +103,16 @@ def decode_b01_map_payload(raw_payload: bytes, *, local_key: str, serial: str, m
         for candidate in candidates:
             variants = [candidate]
             try:
-                text = candidate.decode("ascii").strip()
-                if len(text) > 16 and all(char in "0123456789abcdefABCDEF" for char in text[:32]):
+                text = candidate.strip().decode("ascii")
+            except UnicodeDecodeError:
+                text = ""
+            if text and len(text) >= 32 and len(text) % 2 == 0 and all(
+                char in "0123456789abcdefABCDEF" for char in text
+            ):
+                try:
                     variants.append(bytes.fromhex(text))
-            except Exception:
-                pass
+                except ValueError:
+                    pass
             for variant in variants:
                 try:
                     return zlib.decompress(variant)

tests/protocols/test_b01_map_protocol.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import base64
+import gzip
 import zlib
 from pathlib import Path
 
@@ -11,14 +12,14 @@
 
 from roborock.protocols.b01_map_protocol import decode_b01_map_payload, derive_map_key
 
-FIXTURE = Path(__file__).resolve().parents[1] / "map" / "testdata" / "raw-mqtt-map301.bin.inflated.bin"
+FIXTURE = Path(__file__).resolve().parents[1] / "map" / "testdata" / "raw-mqtt-map301.bin.inflated.bin.gz"
 
 
 def test_decode_b01_map_payload_round_trip() -> None:
     local_key = "abcdefghijklmnop"
     serial = "testsn012345"
     model = "roborock.vacuum.sc05"
-    inflated = FIXTURE.read_bytes()
+    inflated = gzip.decompress(FIXTURE.read_bytes())
 
     compressed = zlib.compress(inflated)
     map_key = derive_map_key(serial, model)