287 lines
12 KiB
Python
287 lines
12 KiB
Python
"""Firmware images: validation, integrity and layout.
|
|
|
|
Two origins:
|
|
* factory - bundled with the tool (firmware/firmware.toml), SHA-256 locked;
|
|
* developer - a newer build the firmware team sends to technicians,
|
|
imported into a per-user library (~/KotaTechnicianTool/firmware).
|
|
Two kinds:
|
|
* merged - bootloader at 0x0 + partition table + app (like new_system.bin);
|
|
App-only or Full flashing;
|
|
* app - the application binary only (build/kota-sensorboard.bin);
|
|
App-only flashing, laid out like the factory image.
|
|
|
|
Every image is checked structurally (ESP32-S3, kota-sensorboard app
|
|
descriptor, segment checksum and the appended SHA-256) so a truncated or
|
|
corrupted file is refused before anything is written.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import datetime as dt
|
|
import hashlib
|
|
import json
|
|
import shutil
|
|
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"
|
|
LIBRARY_DIR = Path.home() / "KotaTechnicianTool" / "firmware"
|
|
LIBRARY_INDEX = "library.json"
|
|
|
|
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"}
|
|
_MONTHS = {m: i for i, m in enumerate(
|
|
["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"], 1)}
|
|
|
|
|
|
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 # "Jun 29 2026 11:03:32" as in the app descriptor
|
|
idf: str
|
|
elf_sha: str # hex
|
|
|
|
@property
|
|
def built_at(self) -> dt.datetime | None:
|
|
try:
|
|
mon, day, year, hms = self.built.split()
|
|
h, m, s = (int(x) for x in hms.split(":"))
|
|
return dt.datetime(int(year), _MONTHS[mon], int(day), h, m, s)
|
|
except (ValueError, KeyError):
|
|
return None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FirmwareImage:
|
|
path: Path
|
|
name: str
|
|
notes: str
|
|
origin: str # "factory" | "developer"
|
|
kind: str # "merged" | "app"
|
|
sha256: str
|
|
flash_size: str
|
|
partitions: tuple[Partition, ...]
|
|
partition_table: bytes # expected table on the box
|
|
app_offset: int # where the app starts inside the file
|
|
app_length: int # exact app image length (from its structure)
|
|
app: AppInfo
|
|
|
|
@property
|
|
def data(self) -> bytes:
|
|
return self.path.read_bytes()
|
|
|
|
@property
|
|
def is_developer(self) -> bool:
|
|
return self.origin == "developer"
|
|
|
|
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")
|
|
|
|
def app_bytes(self) -> bytes:
|
|
return self.data[self.app_offset:self.app_offset + self.app_length]
|
|
|
|
def blank_otadata(self) -> bytes:
|
|
"""Erased otadata: the bootloader then boots ota_0."""
|
|
return b"\xff" * self.partition("otadata").size
|
|
|
|
def label(self) -> str:
|
|
tag = "factory" if self.origin == "factory" else "developer build"
|
|
return f"{self.app.version} built {self.app.built} [{self.app.elf_sha[:9]}] - {tag}"
|
|
|
|
|
|
# ---------------------------------------------------------------- parsing
|
|
|
|
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_image_length(data: bytes, offset: int = 0) -> int:
|
|
"""Validate an ESP app image at `offset`; return its exact length.
|
|
|
|
Checks the segment checksum and, when the header says so, the appended
|
|
SHA-256 - the same checks the bootloader performs before booting it.
|
|
"""
|
|
if len(data) < offset + 24 or data[offset] != IMAGE_MAGIC:
|
|
raise ImageError(f"no firmware image at 0x{offset:x}")
|
|
segments = data[offset + 1]
|
|
chip_id = struct.unpack_from("<H", data, offset + 12)[0]
|
|
if chip_id != ESP32S3_CHIP_ID:
|
|
raise ImageError(f"image is for chip id {chip_id}, not ESP32-S3")
|
|
hash_appended = data[offset + 23] == 1
|
|
pos, checksum = offset + 24, 0xEF
|
|
for _ in range(segments):
|
|
if pos + 8 > len(data):
|
|
raise ImageError("image is truncated (segment header missing)")
|
|
_, seg_len = struct.unpack_from("<II", data, pos)
|
|
pos += 8
|
|
if pos + seg_len > len(data):
|
|
raise ImageError("image is truncated (segment data missing)")
|
|
for b in data[pos:pos + seg_len]:
|
|
checksum ^= b
|
|
pos += seg_len
|
|
pos += 15 - ((pos - offset) % 16) # pad so the checksum ends a 16-byte block
|
|
if pos >= len(data) or data[pos] != checksum:
|
|
raise ImageError("image checksum is wrong - the file is corrupted")
|
|
pos += 1
|
|
if hash_appended:
|
|
if pos + 32 > len(data):
|
|
raise ImageError("image is truncated (SHA-256 missing)")
|
|
if hashlib.sha256(data[offset:pos]).digest() != data[pos:pos + 32]:
|
|
raise ImageError("image SHA-256 does not match its content - the file is corrupted")
|
|
pos += 32
|
|
return pos - offset
|
|
|
|
|
|
def _app_info(data: bytes, offset: int) -> AppInfo:
|
|
d = data[offset + 0x20:offset + 0x20 + 256]
|
|
if len(d) < 256 or struct.unpack_from("<I", d)[0] != APP_DESC_MAGIC:
|
|
raise ImageError("firmware has no ESP-IDF app descriptor")
|
|
s = lambda o, n: d[o:o + n].split(b"\0")[0].decode("ascii", "replace")
|
|
info = 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())
|
|
if info.project != "kota-sensorboard":
|
|
raise ImageError(f"firmware is {info.project!r}, not kota-sensorboard")
|
|
return info
|
|
|
|
|
|
def inspect_image(path: Path, name: str = "", notes: str = "", origin: str = "factory",
|
|
reference: FirmwareImage | None = None) -> FirmwareImage:
|
|
"""Validate a merged or app-only image. App-only images need a `reference`
|
|
(the factory image) that provides the flash layout."""
|
|
data = path.read_bytes()
|
|
if len(data) >= PT_OFFSET + PT_SIZE and data[0] == IMAGE_MAGIC and data[PT_OFFSET:PT_OFFSET + 2] == b"\xaa\x50":
|
|
app_image_length(data, 0) # bootloader must be intact
|
|
table = data[PT_OFFSET:PT_OFFSET + PT_SIZE]
|
|
parts = parse_partition_table(table)
|
|
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")
|
|
kind, app_off, flash_size = "merged", ota0.offset, FLASH_SIZES.get(data[3] >> 4, "?")
|
|
elif data[:1] == bytes([IMAGE_MAGIC]):
|
|
if reference is None:
|
|
raise ImageError("app-only image needs the factory image as layout reference")
|
|
kind, app_off = "app", 0
|
|
parts, table, flash_size = reference.partitions, reference.partition_table, reference.flash_size
|
|
ota0 = next(p for p in parts if p.label == "ota_0")
|
|
else:
|
|
raise ImageError("not an ESP32 firmware file")
|
|
length = app_image_length(data, app_off)
|
|
if length > ota0.size:
|
|
raise ImageError(f"firmware ({length} bytes) does not fit the ota_0 partition ({ota0.size} bytes)")
|
|
app = _app_info(data, app_off)
|
|
return FirmwareImage(path, name or path.name, notes, origin, kind, hashlib.sha256(data).hexdigest(),
|
|
flash_size, parts, table, app_off, length, app)
|
|
|
|
|
|
# ------------------------------------------------------------- catalogue
|
|
|
|
def factory_images(manifest: Path = MANIFEST) -> list[FirmwareImage]:
|
|
"""Bundled images whose file exists and matches the SHA-256 in the manifest."""
|
|
if not manifest.exists():
|
|
return []
|
|
out = []
|
|
for e in tomllib.loads(manifest.read_text(encoding="utf-8")).get("image", []):
|
|
path = manifest.parent / e["file"]
|
|
if not path.exists():
|
|
continue
|
|
img = inspect_image(path, e.get("name", ""), e.get("notes", ""), "factory")
|
|
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
|
|
|
|
|
|
class FirmwareLibrary:
|
|
"""Developer builds imported by the technician (persist across tool versions)."""
|
|
|
|
def __init__(self, root: Path = LIBRARY_DIR):
|
|
self.root = root
|
|
|
|
def _index(self) -> list[dict]:
|
|
f = self.root / LIBRARY_INDEX
|
|
return json.loads(f.read_text(encoding="utf-8")) if f.exists() else []
|
|
|
|
def _save(self, entries: list[dict]) -> None:
|
|
self.root.mkdir(parents=True, exist_ok=True)
|
|
(self.root / LIBRARY_INDEX).write_text(json.dumps(entries, indent=2), encoding="utf-8")
|
|
|
|
def images(self, reference: FirmwareImage | None) -> tuple[list[FirmwareImage], list[str]]:
|
|
"""(valid images, problems) - a changed or broken file is reported, not offered."""
|
|
imgs, problems = [], []
|
|
for e in self._index():
|
|
path = self.root / e["file"]
|
|
try:
|
|
img = inspect_image(path, e.get("name", ""), e.get("notes", ""), "developer", reference)
|
|
if img.sha256 != e["sha256"]:
|
|
raise ImageError("file changed since it was imported")
|
|
imgs.append(img)
|
|
except (OSError, ImageError) as exc:
|
|
problems.append(f"{e.get('original', e['file'])}: {exc}")
|
|
return imgs, problems
|
|
|
|
def import_file(self, src: Path, reference: FirmwareImage | None, notes: str = "") -> FirmwareImage:
|
|
img = inspect_image(src, origin="developer", reference=reference) # validate before copying
|
|
entries = self._index()
|
|
existing = next((e for e in entries if e["sha256"] == img.sha256), None)
|
|
if existing:
|
|
return inspect_image(self.root / existing["file"], existing["name"], existing["notes"],
|
|
"developer", reference)
|
|
self.root.mkdir(parents=True, exist_ok=True)
|
|
fname = f"{img.app.elf_sha[:9]}_{img.sha256[:8]}.bin"
|
|
shutil.copyfile(src, self.root / fname)
|
|
name = f"Developer build {img.app.version} ({img.app.built})"
|
|
entries.append({"file": fname, "sha256": img.sha256, "name": name, "original": src.name,
|
|
"imported": dt.datetime.now().isoformat(timespec="seconds"), "notes": notes})
|
|
self._save(entries)
|
|
return inspect_image(self.root / fname, name, notes, "developer", reference)
|
|
|
|
def remove(self, sha256: str) -> None:
|
|
entries = self._index()
|
|
keep = [e for e in entries if e["sha256"] != sha256]
|
|
for e in entries:
|
|
if e["sha256"] == sha256:
|
|
(self.root / e["file"]).unlink(missing_ok=True)
|
|
self._save(keep)
|
|
|
|
|
|
# Backwards-compatible name used by earlier callers/tests.
|
|
official_images = factory_images
|