Do not crash after KVS unpack failure

tuya-cloudcutter · Dec 1, 2023 · d9f6b14 · d9f6b14
1 parent 665743e
commit d9f6b14
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/bk7231tools/__main__.py b/bk7231tools/__main__.py
@@ -303,15 +303,25 @@ def dissect_dump_file(args):
             break
         pos, storage_data = result
         aes = AES.new(key=KEY_MASTER, mode=AES.MODE_ECB)
-        kvs = KVStorage.unpack(storage_data, aes=aes)
+        try:
+            kvs = KVStorage.unpack(storage_data, aes=aes)
+        except Exception:
+            print("!!! Couldn't unpack KVStorage, perhaps dump file is invalid?")
+            traceback.print_exc()
+            break
         keys = list(kvs.indexes.keys())
         print(f"\t{pos:#06x}: {kvs.length // 1024:d} KiB - {len(keys)} keys")
         print("\n".join(f"\t- '{key}'" for key in keys))
         if args.extract:
             dumpfile_name = Path(dumpfile).stem
 
             out_name = os.path.join(output_directory, f"{dumpfile_name}_storage.json")
-            kvs_data = kvs.read_all_values_parsed()
+            try:
+                kvs_data = kvs.read_all_values_parsed()
+            except Exception:
+                print("!!! Couldn't read parsed KVS values")
+                traceback.print_exc()
+                break
             with open(out_name, "w") as f:
                 print(f"\t\textracted all keys to {out_name}")
                 json.dump(kvs_data, f, indent="\t")