Technician tool v0.4

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

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