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

kota_tech/flash/plan.py

@@ -0,0 +1,101 @@
+"""Flash plans: which esptool steps run, in which order.
+
+App-only (default, repairs): check the box has the same partition table,
+then write the app to ota_0 and erase otadata so ota_0 boots. NVS is kept:
+Matter pairing, OTA registration and Wi-Fi survive.
+
+Full (new or wiped boards): write the whole factory image at 0x0, the same
+as the factory flash_all.bat. This erases NVS: pairing and OTA registration
+are lost.
+"""
+
+from __future__ import annotations
+
+import enum
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from .image import PT_OFFSET, PT_SIZE, FirmwareImage
+
+BAUD = "460800"
+
+
+class Mode(str, enum.Enum):
+    APP_ONLY = "app_only"
+    FULL = "full"
+
+
+@dataclass
+class Step:
+    kind: str                  # "backup" | "read_pt" | "write"
+    title: str
+    args: list[str]            # esptool arguments (no interpreter prefix)
+    est_s: int                 # rough duration for the UI
+    output: Path | None = None # file produced by the step (backup / table)
+
+
+@dataclass
+class Plan:
+    mode: Mode
+    image: FirmwareImage
+    port: str
+    steps: list[Step] = field(default_factory=list)
+    workdir: Path | None = None
+
+    def summary(self) -> list[str]:
+        return [f"{i}. {s.title} (~{_dur(s.est_s)})" for i, s in enumerate(self.steps, 1)]
+
+
+def _dur(s: int) -> str:
+    return f"{s // 60} min" if s >= 90 else f"{s} s"
+
+
+def _common(port: str, after: str) -> list[str]:
+    return ["--chip", "esp32s3", "--port", port, "--baud", BAUD,
+            "--before", "default_reset", "--after", after]
+
+
+def build_plan(image: FirmwareImage, mode: Mode, port: str, workdir: Path,
+               backup_to: Path | None = None) -> Plan:
+    workdir.mkdir(parents=True, exist_ok=True)
+    plan = Plan(mode, image, port, workdir=workdir)
+    if backup_to is not None:
+        size = {"4MB": "0x400000", "8MB": "0x800000", "16MB": "0x1000000"}.get(image.flash_size, "0x800000")
+        plan.steps.append(Step(
+            "backup", f"Back up the box's complete flash to {backup_to.name}",
+            _common(port, "no_reset") + ["read_flash", "0", size, str(backup_to)],
+            est_s=200 if size == "0x800000" else 100, output=backup_to))
+
+    write_opts = ["write_flash", "--flash_mode", "keep", "--flash_freq", "keep", "--flash_size", "keep"]
+    if mode == Mode.APP_ONLY:
+        pt_file = workdir / "box_partition_table.bin"
+        plan.steps.append(Step(
+            "read_pt", "Check the box has the same partition layout",
+            _common(port, "no_reset") + ["read_flash", hex(PT_OFFSET), hex(PT_SIZE), str(pt_file)],
+            est_s=5, output=pt_file))
+        app_file = workdir / "app_ota0.bin"
+        ota_file = workdir / "otadata_blank.bin"
+        app_file.write_bytes(image.app_bytes())
+        ota_file.write_bytes(image.blank_otadata())
+        otadata, ota0 = image.partition("otadata"), image.partition("ota_0")
+        plan.steps.append(Step(
+            "write", "Write the firmware app (keeps pairing, OTA registration, settings)",
+            _common(port, "hard_reset") + write_opts + [
+                hex(otadata.offset), str(ota_file), hex(ota0.offset), str(app_file)],
+            est_s=20))
+    else:
+        plan.steps.append(Step(
+            "write", "Write the complete factory image (erases pairing and OTA registration)",
+            _common(port, "hard_reset") + write_opts + ["0x0", str(image.path)],
+            est_s=25))
+    return plan
+
+
+def check_partition_table(box_table: bytes, image: FirmwareImage) -> str | None:
+    """None if App-only is safe, otherwise the reason it is not."""
+    if box_table.strip(b"\xff") == b"":
+        return "The box has no partition table (blank or erased chip). Use Full mode."
+    if box_table != image.partition_table:
+        return ("The box uses a different partition layout than this image. "
+                "App-only could leave it unbootable - use Full mode.")
+    return None

kota_tech/flash/runner.py

@@ -0,0 +1,118 @@
+"""Runs a flash plan: one esptool subprocess per step, in a background thread.
+
+A subprocess (instead of calling esptool in-process) keeps esptool's
+stdout/sys.exit handling away from the GUI and lets the tool read progress
+line by line. Callbacks are invoked from the worker thread.
+"""
+
+from __future__ import annotations
+
+import re
+import subprocess
+import sys
+import threading
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from .plan import Plan, check_partition_table
+
+_PROGRESS = re.compile(r"\((\d{1,3}) ?%\)")
+
+# esptool output -> technician-facing hint
+HINTS = [
+    (re.compile(r"could not open port|Access is denied|PermissionError|Resource busy", re.I),
+     "The USB port is in use. Close other programs that use the sensorbox port (serial monitors, Arduino IDE) and try again."),
+    (re.compile(r"Failed to connect to ESP32-S3|No serial data received|Wrong boot mode", re.I),
+     "The box did not enter flashing mode by itself. Open the case, hold the BOOT button, press and release RESET, "
+     "release BOOT, then press Start again."),
+    (re.compile(r"device reports readiness to read but returned no data|The port disappeared|device disconnected", re.I),
+     "The USB connection dropped. Use a short data-capable USB cable directly on the laptop (no hub) and try again."),
+    (re.compile(r"Invalid head of packet|timed out", re.I),
+     "Communication errors. Try another USB cable or port."),
+]
+
+
+@dataclass
+class Result:
+    ok: bool
+    message: str
+    hint: str = ""
+    failed_step: int | None = None
+
+
+class FlashRunner:
+    def __init__(self, plan: Plan,
+                 on_step: Callable[[int, str], None],
+                 on_progress: Callable[[int], None],
+                 on_output: Callable[[str], None],
+                 on_done: Callable[[Result], None],
+                 esptool_cmd: list[str] | None = None):
+        self.plan = plan
+        self._on_step, self._on_progress = on_step, on_progress
+        self._on_output, self._on_done = on_output, on_done
+        self._cmd = esptool_cmd or [sys.executable, "-m", "esptool"]
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> None:
+        self._thread = threading.Thread(target=self._run, name="flash", daemon=True)
+        self._thread.start()
+
+    def join(self, timeout: float | None = None) -> None:
+        if self._thread:
+            self._thread.join(timeout)
+
+    def _run(self) -> None:
+        try:
+            self._on_done(self._execute())
+        except Exception as exc:  # never leave the UI waiting
+            self._on_done(Result(False, f"Internal error: {exc}"))
+
+    def _execute(self) -> Result:
+        for i, step in enumerate(self.plan.steps):
+            self._on_step(i, step.title)
+            self._on_progress(0)
+            code, output = self._esptool(step.args)
+            if code != 0:
+                return Result(False, f"Step {i + 1} failed: {step.title}", _hint(output), i)
+            if step.kind == "read_pt":
+                problem = check_partition_table(step.output.read_bytes(), self.plan.image)
+                if problem:
+                    return Result(False, problem, failed_step=i)
+            self._on_progress(100)
+        return Result(True, "Firmware written and verified by esptool.")
+
+    def _esptool(self, args: list[str]) -> tuple[int, str]:
+        cmd = self._cmd + args
+        self._on_output("$ esptool " + " ".join(args))
+        flags = subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0
+        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+                                creationflags=flags)
+        collected: list[str] = []
+        buf = b""
+        assert proc.stdout is not None
+        while chunk := proc.stdout.read1(4096):   # returns as soon as bytes arrive
+            buf += chunk
+            # esptool redraws progress with '\r'; treat it as a line break.
+            *lines, buf = re.split(rb"[\r\n]", buf)
+            for raw in lines:
+                self._line(raw, collected)
+        if buf:
+            self._line(buf, collected)
+        return proc.wait(), "\n".join(collected)
+
+    def _line(self, raw: bytes, collected: list[str]) -> None:
+        text = raw.decode("utf-8", "replace").strip()
+        if not text:
+            return
+        collected.append(text)
+        if m := _PROGRESS.search(text):
+            self._on_progress(min(int(m.group(1)), 100))
+            return   # progress lines are not worth listing one by one
+        self._on_output(text)
+
+
+def _hint(output: str) -> str:
+    for pattern, hint in HINTS:
+        if pattern.search(output):
+            return hint
+    return "See the esptool output below. Save it and send it to the firmware team if the problem repeats."