Try to fix decrypted payload quirks if it fails to parse as json (#236)

miio/protocol.py

@@ -160,17 +160,32 @@ def _decode(self, obj, context, path):
             _LOGGER.debug("Unable to decrypt, returning raw bytes: %s", obj)
             return obj
 
-        decoded = decrypted.decode('utf-8')
-        try:
-            return json.loads(decoded)
-        except:
+        # list of adaption functions for malformed json payload (quirks)
+        decrypted_quirks = [
+            # try without modifications first
+            lambda decrypted_bytes: decrypted_bytes,
+            # powerstrip returns malformed JSON if the device is not
+            # connected to the cloud, so we try to fix it here carefully.
+            lambda decrypted_bytes: decrypted_bytes.replace(b',,"otu_stat"', b',"otu_stat"'),
+            # xiaomi cloud returns malformed json when answering _sync.batch_gen_room_up_url
+            # command so try to sanitize it
+            lambda decrypted_bytes:
+                decrypted_bytes[:decrypted_bytes.rfind(b'\x00')]
+                if b'\x00' in decrypted_bytes
+                else decrypted_bytes
+        ]
+
+        for i, quirk in enumerate(decrypted_quirks):
+            decoded = quirk(decrypted).decode('utf-8')
             try:
-                # powerstrip returns invalid JSON if the device is not
-                # connected to the cloud, so we try to fix it here carefully.
-                decoded = decoded.replace(',,"otu_stat"', ',"otu_stat"')
                 return json.loads(decoded)
             except Exception as ex:
-                _LOGGER.error("unable to parse json '%s': %s", decoded, ex)
+                # log the error when decrypted bytes couldn't be loaded
+                # after trying all quirk adaptions
+                if i == len(decrypted_quirks) - 1:
+                    _LOGGER.error("unable to parse json '%s': %s", decoded, ex)
+
+        return None
 
 
 Message = Struct(

miio/tests/test_protocol.py

@@ -1,5 +1,7 @@
 from unittest import TestCase
 from .. import Utils
+from ..protocol import Message
+import binascii
 
 
 class TestProtocol(TestCase):
@@ -17,7 +19,7 @@ def test_encrypt(self):
 
         encrypted = Utils.encrypt(payload, token)
         decrypted = Utils.decrypt(encrypted, token)
-        self.assertEquals(payload, decrypted)
+        assert payload == decrypted
 
     def test_invalid_token(self):
         payload = b"hello world"
@@ -32,3 +34,45 @@ def test_invalid_token(self):
             Utils.encrypt(payload, wrong_length)
         with self.assertRaises(ValueError):
             Utils.decrypt(payload, wrong_length)
+
+    def test_decode_json_payload(self):
+        token = bytes.fromhex(32 * '0')
+        ctx = {'token': token}
+
+        def build_msg(data):
+            encrypted_data = Utils.encrypt(data, token)
+
+            # header
+            magic = binascii.unhexlify(b'2131')
+            length = (32 + len(encrypted_data)).to_bytes(2, byteorder='big')
+            unknown = binascii.unhexlify(b'00000000')
+            did = binascii.unhexlify(b'01234567')
+            epoch = binascii.unhexlify(b'00000000')
+
+            checksum = Utils.md5(magic+length+unknown+did+epoch+token+encrypted_data)
+
+            return magic+length+unknown+did+epoch+checksum+encrypted_data
+
+        # can parse message with valid json
+        serialized_msg = build_msg(b'{"id": 123456}')
+        parsed_msg = Message.parse(serialized_msg, **ctx)
+        assert parsed_msg.data.value
+        assert isinstance(parsed_msg.data.value, dict)
+        assert parsed_msg.data.value['id'] == 123456
+
+        # can parse message with invalid json for edge case powerstrip
+        # when not connected to cloud
+        serialized_msg = build_msg(b'{"id": 123456,,"otu_stat":0}')
+        parsed_msg = Message.parse(serialized_msg, **ctx)
+        assert parsed_msg.data.value
+        assert isinstance(parsed_msg.data.value, dict)
+        assert parsed_msg.data.value['id'] == 123456
+        assert parsed_msg.data.value['otu_stat'] == 0
+
+        # can parse message with invalid json for edge case xiaomi cloud
+        # reply to _sync.batch_gen_room_up_url
+        serialized_msg = build_msg(b'{"id": 123456}\x00k')
+        parsed_msg = Message.parse(serialized_msg, **ctx)
+        assert parsed_msg.data.value
+        assert isinstance(parsed_msg.data.value, dict)
+        assert parsed_msg.data.value['id'] == 123456