ha-ninstrip/custom_components/ninstrip/firmware.py

"""OTA-Abbilder fuer ninstrip: Anhaenger, Versionsreihenfolge, WLAN-Pruefung, Katalog.

Bewusst frei von Home-Assistant-Importen. Blockierende Funktionen sind
benannt; die Aufrufer schicken sie in einen Executor.

Der Anhaenger ist derselbe wie bei ninlight (512 Byte, Magie "NINL") - genau
deshalb kann die Migrationsstation von dort auch ninstrip-Abbilder ausliefern.
Ein Unterschied bleibt: ninlight laeuft auf dem W600 und traegt seine
Nutzdaten gzip-verpackt, der ESP32 nicht. Die WLAN-Pruefung sucht hier also
im Klartext (siehe `contains_ssid`).
"""

from __future__ import annotations

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

_LOGGER = logging.getLogger(__name__)

TRAILER_SIZE = 512
TRAILER_MAGIC = b"NINL"
TRAILER_FORMAT = 1
# tools/mkimage.py: device[16] und fw_version[16], jeweils NUL-abgeschlossen.
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):
    """Die Datei ist kein gueltiges ninstrip-OTA-Abbild."""


@dataclass(frozen=True, slots=True)
class FirmwareVersion:
    """Eine gelesene `git describe`-Version. Der Hash zaehlt bewusst nicht mit."""

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

    @property
    def sort_key(self) -> tuple[int, int, int, int, int, int]:
        # Ein Release Candidate kommt vor der Freigabe derselben Grundversion.
        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:
        # Die Kuerzung kann mitten in der Commit- oder rc-Zahl gelandet sein:
        # dann ist die Version nicht vergleichbar.
        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]:
    """(device_type, version) eines gueltigen Abbildes, sonst InvalidImageError."""
    if len(data) <= TRAILER_SIZE:
        raise InvalidImageError("zu kurz fuer einen Anhaenger")
    payload, trailer = data[:-TRAILER_SIZE], data[-TRAILER_SIZE:]
    if trailer[0:4] != TRAILER_MAGIC:
        raise InvalidImageError("kein NINL-Anhaenger")
    if trailer[4] != TRAILER_FORMAT:
        raise InvalidImageError(f"unbekanntes Anhaengerformat {trailer[4]}")
    if trailer[37:69] != hashlib.sha256(payload).digest():
        raise InvalidImageError("sha256 passt nicht")
    device_type = _text_field(trailer[5:21])
    version = _text_field(trailer[21:37])
    if not device_type or not version:
        raise InvalidImageError("Geraetetyp oder Version leer")
    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:
    """Wahr, wenn das Abbild `ssid` als NUL-gerahmte C-Zeichenkette traegt.

    Anders als bei ninlight wird nichts ausgepackt: ein ESP32-Abbild liegt
    unverpackt vor, die einkompilierte SSID steht als Zeichenkette drin. Am
    Geraet gegengeprueft - in einem mit sdkconfig.local gebauten Abbild kommt
    sie genau einmal NUL-gerahmt vor.
    """
    if not ssid or len(data) <= TRAILER_SIZE:
        return False
    return b"\x00" + ssid.encode() + b"\x00" in data[:-TRAILER_SIZE]


@dataclass(frozen=True, slots=True)
class FirmwareImage:
    """Ein gueltiges Abbild im Firmware-Ordner."""

    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:
    """Alle gueltigen Abbilder eines Ordners. Nicht threadsicher: Aufrufer ordnen die Laeufe."""

    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() laeuft in einem Executor-Thread, waehrend contains_ssid() in
        # einem anderen laufen kann (etwa die WLAN-Gegenprobe einer laufenden
        # Aktualisierung gegen einen turnusmaessigen Durchlauf): beide lesen
        # und aendern _ssid_hits.
        self._ssid_hits_lock = threading.Lock()

    def scan(self) -> tuple[FirmwareImage, ...]:
        """Blockierend. Ungueltige Dateien werden je (Pfad, mtime, Groesse) einmal gemeldet."""
        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("Uebergehe Firmware-Datei %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]:
        """Abbilder genau dieser Platine, die neuer sind als `installed` - neuestes zuerst."""
        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_and_read(self, image: FirmwareImage) -> bytes | None:
        """Blockierend. Liefert genau die geprueften Bytes, sonst None.

        Wer mit dem Inhalt weiterarbeitet (ausliefern, nach der SSID
        durchsuchen), muss diese Bytes nehmen und die Datei nicht erneut
        lesen - dazwischen kann sie sich aendern.
        """
        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:
        """Blockierend, gemerkt je (Datei, 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