Technician tool v0.4

Imported from the Google Drive folder 004-technician_tool/v0.4.

kota_tech/flash/image.py

@@ -0,0 +1,152 @@
+"""Official firmware images: validation and layout.
+
+Only images listed in firmware/firmware.toml with a matching SHA-256 are
+offered for flashing: boxes normally get new firmware over OTA, so the tool
+only ever writes the official factory image.
+
+The factory image is a merged ESP32-S3 image (bootloader at 0x0, partition
+table at 0x8000, app at ota_0). App-only flashing takes the app out of it
+and writes it to ota_0; the layout comes from the image's own partition
+table, not from constants.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import struct
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+
+FIRMWARE_DIR = Path(__file__).resolve().parents[2] / "firmware"
+MANIFEST = FIRMWARE_DIR / "firmware.toml"
+
+PT_OFFSET = 0x8000
+PT_SIZE = 0xC00
+ESP32S3_CHIP_ID = 9
+IMAGE_MAGIC = 0xE9
+APP_DESC_MAGIC = 0xABCD5432
+FLASH_SIZES = {0: "1MB", 1: "2MB", 2: "4MB", 3: "8MB", 4: "16MB", 5: "32MB"}
+
+
+class ImageError(ValueError):
+    pass
+
+
+@dataclass(frozen=True)
+class Partition:
+    label: str
+    type: int
+    subtype: int
+    offset: int
+    size: int
+
+
+@dataclass(frozen=True)
+class AppInfo:
+    project: str
+    version: str
+    built: str
+    idf: str
+    elf_sha: str       # hex
+
+
+@dataclass(frozen=True)
+class FirmwareImage:
+    path: Path
+    name: str
+    notes: str
+    sha256: str
+    flash_size: str
+    partitions: tuple[Partition, ...]
+    app: AppInfo
+
+    @property
+    def data(self) -> bytes:
+        return self.path.read_bytes()
+
+    def partition(self, label: str) -> Partition:
+        for p in self.partitions:
+            if p.label == label:
+                return p
+        raise ImageError(f"partition {label!r} not in image")
+
+    @property
+    def partition_table(self) -> bytes:
+        return self.data[PT_OFFSET:PT_OFFSET + PT_SIZE]
+
+    def app_bytes(self) -> bytes:
+        """The application image stored in the ota_0 region of the merged image."""
+        # No stripping of trailing 0xFF: the image ends in a SHA-256 digest
+        # that may legitimately end in 0xFF. The merged file ends with the app.
+        ota0 = self.partition("ota_0")
+        return self.data[ota0.offset:ota0.offset + ota0.size]
+
+    def blank_otadata(self) -> bytes:
+        """Erased otadata: the bootloader then boots ota_0."""
+        return b"\xff" * self.partition("otadata").size
+
+
+def parse_partition_table(raw: bytes) -> tuple[Partition, ...]:
+    parts = []
+    for i in range(0, len(raw) - 31, 32):
+        entry = raw[i:i + 32]
+        magic = entry[:2]
+        if magic == b"\xaa\x50":
+            _, ptype, sub, off, size = struct.unpack_from("<HBBII", entry)
+            label = entry[12:28].split(b"\0")[0].decode("ascii", "replace")
+            parts.append(Partition(label, ptype, sub, off, size))
+        elif magic in (b"\xeb\xeb", b"\xff\xff"):   # MD5 entry / end
+            break
+        else:
+            raise ImageError(f"corrupt partition table entry at 0x{PT_OFFSET + i:x}")
+    return tuple(parts)
+
+
+def _app_info(data: bytes, offset: int) -> AppInfo:
+    if len(data) < offset + 0x20 + 256 or data[offset] != IMAGE_MAGIC:
+        raise ImageError(f"no application image at 0x{offset:x}")
+    d = data[offset + 0x20:offset + 0x20 + 256]
+    if struct.unpack_from("<I", d)[0] != APP_DESC_MAGIC:
+        raise ImageError("application has no ESP-IDF app descriptor")
+    s = lambda o, n: d[o:o + n].split(b"\0")[0].decode("ascii", "replace")
+    return AppInfo(project=s(48, 32), version=s(16, 32), built=f"{s(96, 16)} {s(80, 16)}",
+                   idf=s(112, 32), elf_sha=d[144:176].hex())
+
+
+def inspect_image(path: Path, name: str = "", notes: str = "") -> FirmwareImage:
+    data = path.read_bytes()
+    if len(data) < PT_OFFSET + PT_SIZE or data[0] != IMAGE_MAGIC:
+        raise ImageError("not a merged ESP32 image (no bootloader at offset 0)")
+    chip_id = struct.unpack_from("<H", data, 12)[0]
+    if chip_id != ESP32S3_CHIP_ID:
+        raise ImageError(f"image is for chip id {chip_id}, not ESP32-S3")
+    flash_size = FLASH_SIZES.get(data[3] >> 4, "?")
+    parts = parse_partition_table(data[PT_OFFSET:PT_OFFSET + PT_SIZE])
+    labels = {p.label for p in parts}
+    for need in ("nvs", "otadata", "ota_0"):
+        if need not in labels:
+            raise ImageError(f"partition table has no {need!r} partition")
+    ota0 = next(p for p in parts if p.label == "ota_0")
+    app = _app_info(data, ota0.offset)
+    if app.project != "kota-sensorboard":
+        raise ImageError(f"image is {app.project!r}, not kota-sensorboard")
+    return FirmwareImage(path, name or path.name, notes, hashlib.sha256(data).hexdigest(),
+                         flash_size, parts, app)
+
+
+def official_images(manifest: Path = MANIFEST) -> list[FirmwareImage]:
+    """Images from the manifest whose file exists and matches its SHA-256."""
+    if not manifest.exists():
+        return []
+    entries = tomllib.loads(manifest.read_text(encoding="utf-8")).get("image", [])
+    out = []
+    for e in entries:
+        path = manifest.parent / e["file"]
+        if not path.exists():
+            continue
+        img = inspect_image(path, e.get("name", ""), e.get("notes", ""))
+        if img.sha256 != e["sha256"].lower():
+            raise ImageError(f"{path.name}: SHA-256 does not match the manifest - file changed or corrupt")
+        out.append(img)
+    return out