ha-ninlight/custom_components/ninlight/firmware.py

"""Firmware images for ninlight lamps: trailer, version order, Wi-Fi check, catalog.

Deliberately free of Home Assistant imports. Blocking functions are marked; callers run them in an
executor. Facts behind the constants: spec §2 (layout checked against a real image).
"""

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import logging
from pathlib import Path
import re
import threading
import zlib

_LOGGER = logging.getLogger(__name__)

TRAILER_SIZE = 512
TRAILER_MAGIC = b"NINL"
TRAILER_FORMAT = 1
W600_HEADER_SIZE = 56
GZIP_MAGIC = b"\x1f\x8b\x08"
# cmake/write_version.cmake cuts the version hard to 15 characters.
VERSION_FIELD_MAX = 15

_VERSION_RE = re.compile(
    r"^(?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)"
    r"(?:-rc(?P<rc>\d+))?"
    r"(?:-(?P<commits>\d+)(?P<hash>-g[0-9a-f]*)?)?"
    r"(?:-d(?:i(?:r(?:t(?:y)?)?)?)?)?$"
)


class InvalidImageError(ValueError):
    """The file is not a valid ninlight OTA image."""


@dataclass(frozen=True, slots=True)
class FirmwareVersion:
    """A parsed `git describe` version. The hash is ignored on purpose (usually truncated)."""

    base: tuple[int, int, int]
    rc: int | None
    commits: int

    @property
    def sort_key(self) -> tuple[int, int, int, int, int, int]:
        # A release candidate sorts before the release of the same base.
        rc_rank, rc_number = (0, self.rc) if self.rc is not None else (1, 0)
        return (*self.base, rc_rank, rc_number, self.commits)


def parse_version(text: str) -> FirmwareVersion | None:
    match = _VERSION_RE.match(text)
    if match is None:
        return None
    if len(text) == VERSION_FIELD_MAX and text[-1].isdigit() and match["hash"] is None:
        # The cut may have landed inside the commit (or rc) number: not comparable.
        return None
    return FirmwareVersion(
        (int(match["major"]), int(match["minor"]), int(match["patch"])),
        int(match["rc"]) if match["rc"] is not None else None,
        int(match["commits"]) if match["commits"] is not None else 0,
    )


def is_newer(candidate: str, installed: str) -> bool:
    new, old = parse_version(candidate), parse_version(installed)
    if new is None or old is None:
        return False
    return new.sort_key > old.sort_key


def parse_trailer(data: bytes) -> tuple[str, str]:
    """Return (device_type, version) of a valid image or raise InvalidImageError."""
    if len(data) <= TRAILER_SIZE:
        raise InvalidImageError("too short for a trailer")
    payload, trailer = data[:-TRAILER_SIZE], data[-TRAILER_SIZE:]
    if trailer[0:4] != TRAILER_MAGIC:
        raise InvalidImageError("no NINL trailer")
    if trailer[4] != TRAILER_FORMAT:
        raise InvalidImageError(f"unsupported trailer format {trailer[4]}")
    if trailer[37:69] != hashlib.sha256(payload).digest():
        raise InvalidImageError("sha256 mismatch")
    device_type = _text_field(trailer[5:21])
    version = _text_field(trailer[21:37])
    if not device_type or not version:
        raise InvalidImageError("empty device type or version")
    return device_type, version


def _text_field(raw: bytes) -> str:
    return raw.split(b"\x00", 1)[0].decode("ascii", errors="replace")


def contains_ssid(data: bytes, ssid: str) -> bool:
    """True if the decompressed firmware carries `ssid` as a NUL-framed C string (spec §7)."""
    if not ssid or len(data) <= TRAILER_SIZE:
        return False
    payload = data[:-TRAILER_SIZE]
    start = payload.find(GZIP_MAGIC, W600_HEADER_SIZE)
    if start < 0:
        return False
    try:
        raw = zlib.decompressobj(31).decompress(payload[start:])
    except zlib.error:
        return False
    return b"\x00" + ssid.encode() + b"\x00" in raw


@dataclass(frozen=True, slots=True)
class FirmwareImage:
    """One valid image in the firmware folder."""

    path: Path
    device_type: str
    version: str
    size: int
    mtime_ns: int


def _sort_key(image: FirmwareImage) -> tuple[int, ...]:
    version = parse_version(image.version)
    return version.sort_key if version is not None else ()


class FirmwareCatalog:
    """All valid images in one folder. Not thread-safe: callers serialize scans."""

    def __init__(self, directory: Path) -> None:
        self.directory = directory
        self.images: tuple[FirmwareImage, ...] = ()
        self._parsed: dict[tuple[str, int, int], FirmwareImage | None] = {}
        self._ssid_hits: dict[tuple[str, int, int, str], bool] = {}
        # scan() runs in one executor thread while contains_ssid() may run in another (e.g. an
        # install's Wi-Fi re-check racing a periodic rescan): both read and mutate _ssid_hits.
        self._ssid_hits_lock = threading.Lock()

    def scan(self) -> tuple[FirmwareImage, ...]:
        """Blocking. Invalid files are logged once per (path, mtime, size)."""
        self.directory.mkdir(parents=True, exist_ok=True)
        found: list[FirmwareImage] = []
        seen: set[tuple[str, int, int]] = set()
        for path in sorted(self.directory.glob("*.img")):
            try:
                if not path.is_file():
                    continue
                stat = path.stat()
            except OSError:
                continue
            key = (str(path), stat.st_mtime_ns, stat.st_size)
            seen.add(key)
            if key not in self._parsed:
                try:
                    device_type, version = parse_trailer(path.read_bytes())
                except (OSError, InvalidImageError) as err:
                    _LOGGER.warning("Ignoring firmware file %s: %s", path.name, err)
                    self._parsed[key] = None
                else:
                    self._parsed[key] = FirmwareImage(
                        path, device_type, version, stat.st_size, stat.st_mtime_ns
                    )
            image = self._parsed[key]
            if image is not None:
                found.append(image)
        self._parsed = {k: v for k, v in self._parsed.items() if k in seen}
        with self._ssid_hits_lock:
            self._ssid_hits = {k: v for k, v in self._ssid_hits.items() if k[:3] in seen}
        self.images = tuple(found)
        return self.images

    def newer_images(self, device_type: str, installed: str) -> list[FirmwareImage]:
        """Images for exactly this device type that are newer than `installed`, newest first."""
        candidates = [
            image
            for image in self.images
            if image.device_type == device_type and is_newer(image.version, installed)
        ]
        return sorted(candidates, key=_sort_key, reverse=True)

    def verify(self, image: FirmwareImage) -> bool:
        """Blocking. The file is still there, unchanged and valid for the same type/version."""
        try:
            stat = image.path.stat()
            if (stat.st_mtime_ns, stat.st_size) != (image.mtime_ns, image.size):
                return False
            return parse_trailer(image.path.read_bytes()) == (image.device_type, image.version)
        except OSError, InvalidImageError:
            return False

    def verify_and_read(self, image: FirmwareImage) -> bytes | None:
        """Blocking. Like `verify`, but returns the exact bytes that were checked (or None).

        Callers that need to act on the verified content (e.g. serve it, or scan it for an SSID)
        must use these bytes rather than reading the file again afterwards - the file on disk can
        change between the two reads.
        """
        try:
            stat = image.path.stat()
            if (stat.st_mtime_ns, stat.st_size) != (image.mtime_ns, image.size):
                return None
            data = image.path.read_bytes()
        except OSError:
            return None
        try:
            if parse_trailer(data) != (image.device_type, image.version):
                return None
        except InvalidImageError:
            return None
        return data

    def contains_ssid(self, image: FirmwareImage, ssid: str) -> bool:
        """Blocking, cached per (image file, ssid)."""
        key = (str(image.path), image.mtime_ns, image.size, ssid)
        with self._ssid_hits_lock:
            cached = key in self._ssid_hits
            hit = self._ssid_hits.get(key)
        if cached:
            return bool(hit)
        try:
            data = image.path.read_bytes()
        except OSError:
            return False
        hit = contains_ssid(data, ssid)
        with self._ssid_hits_lock:
            self._ssid_hits[key] = hit
        return hit