Reconstruct OTA images from PCAP files (#29)

* [WIP] Reconstruct OTA files out of PCAP captures

* Rename command to `ota reconstruct-from-pcaps`

* Log an error when identical chunks with different values

* Reduce key lengths a little

* Rename `--network-key` to `--add-network-key` and allow specifying multiple

* Always use default network keys

* Print the OTA image info of every dumped image

* Allow adding install codes with `--add-install-code`

* Expand README

Author: puddly

README.md

@@ -175,6 +175,19 @@ $ zigpy ota generate-index --ota-url-root="https://example.org/fw" path/to/firmw
 ...
 ```
 
+## Reconstruct an OTA image from a series of packet captures
+
+Requires the `tshark` binary to be available.
+
+```console
+$ zigpy ota reconstruct-from-pcaps --add-network-key aa:bb:cc:dd:ee:ff:00:11:22:33:44:55:66:77:88:99 --output-root ./extracted/ *.pcap
+Constructing image type=0x298b, version=0x00000005, manuf_code=0x115f: 157424 bytes
+2023-02-22 03:39:51.406 ubuntu zigpy_cli.ota ERROR Missing 48 bytes starting at offset 0x0000ADA0: filling with 0xAB
+2023-02-22 03:39:51.406 ubuntu zigpy_cli.ota ERROR Missing 48 bytes starting at offset 0x000106B0: filling with 0xAB
+Constructing image type=0x298b, version=0x00000009, manuf_code=0x115f: 163136 bytes
+```
+
+
 # PCAP
 ## Re-calculate the FCS on a packet capture
 

zigpy_cli/ota.py

@@ -1,19 +1,35 @@
 from __future__ import annotations
 
+import collections
 import hashlib
 import json
 import logging
 import pathlib
+import subprocess
 
 import click
+import zigpy.types as t
 from zigpy.ota.image import ElementTagId, HueSBLOTAImage, parse_ota_image
 from zigpy.ota.validators import validate_ota_image
+from zigpy.types.named import _hex_string_to_bytes
+from zigpy.util import convert_install_code as zigpy_convert_install_code
 
 from zigpy_cli.cli import cli
+from zigpy_cli.common import HEX_OR_DEC_INT
 
 LOGGER = logging.getLogger(__name__)
 
 
+def convert_install_code(text: str) -> t.KeyData:
+    code = _hex_string_to_bytes(text)
+    key = zigpy_convert_install_code(code)
+
+    if key is None:
+        raise ValueError(f"Invalid install code: {text!r}")
+
+    return key
+
+
 @cli.group()
 def ota():
     pass
@@ -125,3 +141,152 @@ def generate_index(ctx, ota_url_root, output, files):
 
     json.dump(ota_metadata, output, indent=4)
     output.write("\n")
+
+
+@ota.command()
+@click.pass_context
+@click.option(
+    "--add-network-key", "network_keys", type=t.KeyData.convert, multiple=True
+)
+@click.option(
+    "--add-install-code", "install_codes", type=convert_install_code, multiple=True
+)
+@click.option("--fill-byte", type=HEX_OR_DEC_INT, default=0xAB)
+@click.option(
+    "--output-root",
+    type=click.Path(file_okay=False, dir_okay=True, path_type=pathlib.Path),
+    required=True,
+)
+@click.argument("files", nargs=-1, type=pathlib.Path)
+def reconstruct_from_pcaps(
+    ctx, network_keys, install_codes, fill_byte, output_root, files
+):
+    for code in install_codes:
+        print(f"Using key derived from install code: {code}")
+
+    network_keys = (
+        [
+            # ZigBeeAlliance09
+            t.KeyData.convert("5A:69:67:42:65:65:41:6C:6C:69:61:6E:63:65:30:39"),
+            # Z2M default
+            t.KeyData.convert("01:03:05:07:09:0B:0D:0F:00:02:04:06:08:0A:0C:0D"),
+        ]
+        + list(network_keys)
+        + list(install_codes)
+    )
+
+    keys = "\n".join(
+        [f'"{k}","Normal","Network Key {i + 1}"' for i, k in enumerate(network_keys)]
+    )
+    packets = []
+
+    for f in files:
+        proc = subprocess.run(
+            [
+                "tshark",
+                "-o",
+                f"uat:zigbee_pc_keys:{keys}",
+                "-r",
+                str(f),
+                "-T",
+                "json",
+            ],
+            capture_output=True,
+        )
+
+        obj = json.loads(proc.stdout)
+        packets.extend(p["_source"]["layers"] for p in obj)
+
+    ota_sizes = {}
+    ota_chunks = collections.defaultdict(set)
+
+    for packet in packets:
+        if "zbee_zcl" not in packet:
+            continue
+
+        # Ignore non-OTA packets
+        if packet["zbee_aps"]["zbee_aps.cluster"] != "0x0019":
+            continue
+
+        if (
+            packet.get("zbee_zcl", {})
+            .get("Payload", {})
+            .get("zbee_zcl_general.ota.status")
+            == "0x00"
+        ):
+            zcl = packet["zbee_zcl"]["Payload"]
+            image_version = zcl["zbee_zcl_general.ota.file.version"]
+            image_type = zcl["zbee_zcl_general.ota.image.type"]
+            image_manuf_code = zcl["zbee_zcl_general.ota.manufacturer_code"]
+
+            image_key = (image_version, image_type, image_manuf_code)
+
+            if "zbee_zcl_general.ota.image.size" in zcl:
+                image_size = int(zcl["zbee_zcl_general.ota.image.size"])
+                ota_sizes[image_key] = image_size
+            elif "zbee_zcl_general.ota.image.data" in zcl:
+                offset = int(zcl["zbee_zcl_general.ota.file.offset"])
+                data = bytes.fromhex(
+                    zcl["zbee_zcl_general.ota.image.data"].replace(":", "")
+                )
+
+                ota_chunks[image_key].add((offset, data))
+
+    for image_key, image_size in ota_sizes.items():
+        image_version, image_type, image_manuf_code = image_key
+        print(
+            f"Constructing image type={image_type}, version={image_version}"
+            f", manuf_code={image_manuf_code}: {image_size} bytes"
+        )
+
+        buffer = [None] * image_size
+
+        for offset, chunk in sorted(ota_chunks[image_key]):
+            current_value = buffer[offset : offset + len(chunk)]
+
+            if (
+                all(c is not None for c in buffer[offset : offset + len(chunk)])
+                and current_value != chunk
+            ):
+                LOGGER.error(
+                    f"Inconsistent {len(chunk)} bytes starting at offset"
+                    f" 0x{offset:08X}: was {current_value}, now {chunk}"
+                )
+
+            buffer[offset : offset + len(chunk)] = chunk
+
+        missing_indices = [o for o, v in enumerate(buffer) if v is None]
+        missing_ranges = []
+
+        # For readability, combine the list of missing indices into a list of ranges
+        if missing_indices:
+            start = missing_indices[0]
+            count = 0
+
+            for i in missing_indices[1:]:
+                if i == start + count + 1:
+                    count += 1
+                else:
+                    missing_ranges.append((start, count + 1))
+                    start = i
+                    count = 0
+
+            if count > 0:
+                missing_ranges.append((start, count + 1))
+
+        for start, count in missing_ranges:
+            LOGGER.error(
+                f"Missing {count} bytes starting at offset 0x{start:08X}:"
+                f" filling with 0x{fill_byte:02X}"
+            )
+            buffer[start : start + count] = [fill_byte] * count
+
+        filename = output_root / (
+            f"ota_t{image_type}_m{image_manuf_code}_v{image_version}"
+            f"{'_partial' if missing_ranges else ''}.ota"
+        )
+
+        output_root.mkdir(exist_ok=True)
+        filename.write_bytes(bytes(buffer))
+
+        info.callback([filename])