kota-technician-tool/kota_tech/serialio/commands.py

"""Commands to the box over the USB serial port (firmware with usb_console, api=1).

Protocol (kota-sensorboard main/usb_console.hpp): ASCII lines "kota ...",
one reply line each, logged as "KOTA: cmd <name> ok=1|0 [key=value ...]":

    kota ping                            -> cmd ping ok=1 fw=<version> api=1
    kota health                          -> cmd health ok=1 <health fields> | ok=0 err=not_ready
    kota wifi get                        -> cmd wifi_get ok=1 ssid=<hex|-> password_set=0|1 mode=...
    kota wifi set <ssid_hex> <pass_hex>  -> cmd wifi_set ok=1 restart=1, then restart (usb_wifi_set)
    kota wifi clear                      -> cmd wifi_clear ok=1 restart=1, then restart (usb_wifi_clear)

Firmware without the console never answers. The password is only ever sent
hex-encoded and is never written to the log by the box or by this tool.

WifiFlow runs a whole "write / clear Wi-Fi" job as a small state machine
driven by `feed(line)` and `tick(now)`, so it is testable without Qt or a box.
"""

from __future__ import annotations

import enum
import time
from collections.abc import Callable
from dataclasses import dataclass, field

from ..core.logline import LineKind, LogLine

API_VERSION = 1

# Technician-facing text for the firmware's err= codes.
ERRORS = {
    "bad_args": "The box did not understand the command (tool/firmware mismatch).",
    "bad_hex": "The box could not decode the settings (tool/firmware mismatch).",
    "bad_char": "The name or password contains a line break.",
    "ssid_length": "The network name must be 1-32 bytes.",
    "password_length": "The password must be 8-63 characters.",
    "sd_not_mounted": "The box cannot read its SD card, so it cannot store the settings.",
    "write_failed": "Writing ENV.TXT on the SD card failed.",
    "verify_failed": "ENV.TXT did not read back correctly after writing.",
    "unknown_command": "This firmware does not know the command.",
    "not_ready": "The box is still starting; try again in a few seconds.",
}


def hex_utf8(text: str) -> str:
    return text.encode("utf-8").hex()


def unhex_utf8(value: str) -> str | None:
    if not value or value == "-":
        return None
    try:
        return bytes.fromhex(value).decode("utf-8", errors="replace")
    except ValueError:
        return None


def validate_wifi(ssid: str, password: str) -> list[str]:
    """Same limits as the firmware (usb_console.cpp); empty list = OK."""
    problems = []
    if not ssid:
        problems.append("Enter the Wi-Fi network name (SSID).")
    elif len(ssid.encode("utf-8")) > 32:
        problems.append("The network name is longer than 32 bytes.")
    if not 8 <= len(password.encode("utf-8")) <= 63:
        problems.append("The password must be 8-63 characters (WPA2). Open networks are not supported.")
    if any(c in ssid + password for c in "\r\n\0"):
        problems.append("The name or password contains a line break.")
    return problems


def wifi_set_command(ssid: str, password: str) -> str:
    problems = validate_wifi(ssid, password)
    if problems:
        raise ValueError(" ".join(problems))
    return f"kota wifi set {hex_utf8(ssid)} {hex_utf8(password)}"


def loggable(command: str) -> str:
    """The command as it may appear in the session log: never the password."""
    parts = command.split()
    if parts[:3] == ["kota", "wifi", "set"] and len(parts) == 5:
        return f"kota wifi set (network '{unhex_utf8(parts[3])}', password hidden)"
    return command


@dataclass(frozen=True)
class Reply:
    name: str                 # ping | health | wifi_get | wifi_set | wifi_clear | unknown
    ok: bool
    fields: dict[str, str]

    @property
    def error(self) -> str:
        code = self.fields.get("err", "")
        return ERRORS.get(code, f"The box reported an error ({code or 'unknown'}).")


def parse_reply(line: LogLine) -> Reply | None:
    """'I (5) KOTA: cmd wifi_set ok=1 restart=1' -> Reply('wifi_set', True, {...})."""
    if line.kind != LineKind.KOTA or line.record != "cmd" or not line.words:
        return None
    return Reply(line.words[0], line.fields.get("ok") == "1", dict(line.fields))


@dataclass(frozen=True)
class WifiInfo:
    """What `kota wifi get` reported."""
    ssid: str | None
    password_set: bool
    mode: str

    @classmethod
    def from_reply(cls, r: Reply) -> WifiInfo:
        return cls(unhex_utf8(r.fields.get("ssid", "-")), r.fields.get("password_set") == "1",
                   r.fields.get("mode", "?"))

    def describe(self) -> str:
        mode = {"wifi": "local Wi-Fi from the SD card", "matter": "Matter (Wi-Fi given by the Matter app)",
                "provisioning": "Wi-Fi setup portal"}.get(self.mode, self.mode)
        net = f"network '{self.ssid}'" + (", password set" if self.password_set else ", NO password")
        return f"Mode: {mode}. " + (net if self.ssid else "No network stored on the SD card.")


# ------------------------------------------------------------------ flows

class Step(str, enum.Enum):
    PING = "Checking the box accepts USB commands"
    READ = "Reading the Wi-Fi settings"
    SEND = "Storing the settings on the box's SD card"
    RESTART = "Waiting for the box to restart"
    CONNECT = "Waiting for the box to join the Wi-Fi network"
    DONE = "Done"
    FAILED = "Failed"


@dataclass
class FlowEvent:
    step: Step
    text: str
    ok: bool | None = None        # None = in progress, True = finished OK, False = failed


PING_TIMEOUT_S = 3.0
REPLY_TIMEOUT_S = 8.0
RESTART_TIMEOUT_S = 40.0
CONNECT_TIMEOUT_S = 90.0
HEALTH_POLL_S = 5.0


@dataclass
class WifiFlow:
    """One technician job: "read", "set" (ssid+password) or "clear".

    `send(text)` writes one command line to the box (returns False when the
    port is not connected); `on_event` reports progress for the UI.
    """
    action: str
    send: Callable[[str], bool]
    on_event: Callable[[FlowEvent], None]
    ssid: str = ""
    password: str = ""
    now: Callable[[], float] = time.monotonic
    step: Step = Step.PING
    info: WifiInfo | None = None
    _deadline: float = 0.0
    _next_poll: float = 0.0
    _events: list[FlowEvent] = field(default_factory=list)

    @property
    def finished(self) -> bool:
        return self.step in (Step.DONE, Step.FAILED)

    def start(self) -> None:
        if self.action == "set":
            problems = validate_wifi(self.ssid, self.password)
            if problems:
                return self._fail(" ".join(problems))
        self._go(Step.PING, "kota ping", PING_TIMEOUT_S)

    def cancel(self) -> None:
        if not self.finished:
            self._fail("Cancelled.")

    # -------------------------------------------------------- driving
    def feed(self, line: LogLine) -> None:
        if self.finished:
            return
        if self.step == Step.RESTART and line.kind == LineKind.KOTA and line.record == "boot":
            return self._on_boot(line.fields)
        r = parse_reply(line)
        if r is None:
            return
        if self.step == Step.PING and r.name == "ping":
            if not r.ok or int(r.fields.get("api", "0") or 0) < API_VERSION:
                return self._fail("The box answered, but its command version is not supported.")
            self._emit(f"Firmware {r.fields.get('fw', '?')} accepts USB commands.", True)
            if self.action == "read":
                return self._go(Step.READ, "kota wifi get", REPLY_TIMEOUT_S)
            if self.action == "set":
                return self._go(Step.SEND, wifi_set_command(self.ssid, self.password), REPLY_TIMEOUT_S)
            return self._go(Step.SEND, "kota wifi clear", REPLY_TIMEOUT_S)
        if self.step == Step.READ and r.name == "wifi_get":
            if not r.ok:
                return self._fail(r.error)
            self.info = WifiInfo.from_reply(r)
            return self._finish(self.info.describe())
        if self.step == Step.SEND and r.name in ("wifi_set", "wifi_clear"):
            if not r.ok:
                return self._fail(r.error)
            self._emit("The box stored the settings and is restarting.", True)
            self.step, self._deadline = Step.RESTART, self.now() + RESTART_TIMEOUT_S
            return self._emit(Step.RESTART.value + "...")
        if self.step == Step.CONNECT and r.name == "health" and r.ok:
            return self._on_health(r.fields)

    def tick(self) -> None:
        """Call periodically (e.g. every 0.5 s): timeouts and health polling."""
        if self.finished:
            return
        t = self.now()
        if self.step == Step.CONNECT and t >= self._next_poll:
            self._next_poll = t + HEALTH_POLL_S
            self.send("kota health")
        if t < self._deadline:
            return
        if self.step == Step.PING:
            self._fail("The box does not answer USB commands. It needs firmware with the USB console "
                       "(Flash firmware tab), or it is still starting - try again in 10 s.")
        elif self.step == Step.RESTART:
            self._fail("The box stored the settings but no restart was seen. Press Reset & capture in the "
                       "Serial monitor and check the Diagnosis panel.")
        elif self.step == Step.CONNECT:
            self._fail(f"The box did not join '{self.ssid}' within {int(CONNECT_TIMEOUT_S)} s. Check the network "
                       "name and password, that it is a 2.4 GHz network and that it is in range. The settings stay "
                       "on the box; it keeps retrying.")
        else:
            self._fail("The box did not answer. Check the USB connection and try again.")

    # -------------------------------------------------------- internals
    def _on_boot(self, f: dict[str, str]) -> None:
        mode, last = f.get("mode", "?"), f.get("last_restart", "?")
        want = "usb_wifi_set" if self.action == "set" else "usb_wifi_clear"
        if last != want:
            return self._fail(f"The box restarted for another reason ({last}); the settings may not be applied.")
        self._emit(f"Restarted in {mode} mode.", True)
        if self.action == "clear":
            ok = mode == "matter"
            return self._finish("Wi-Fi removed; the box uses Matter again (pair it in the Matter app if needed).") \
                if ok else self._fail(f"The box restarted in {mode} mode instead of Matter.")
        if mode != "wifi":
            return self._fail(f"The box restarted in {mode} mode instead of local Wi-Fi.")
        self.step = Step.CONNECT
        self._deadline = self.now() + CONNECT_TIMEOUT_S
        self._next_poll = self.now() + HEALTH_POLL_S
        self._emit(Step.CONNECT.value + f" '{self.ssid}'...")

    def _on_health(self, f: dict[str, str]) -> None:
        if f.get("wifi") != "connected":
            return
        rssi = f.get("rssi", "?")
        note = " - weak signal, consider moving the box or the router" if rssi.lstrip("-").isdigit() \
            and int(rssi) < -75 else ""
        self._finish(f"Connected to '{self.ssid}' (signal {rssi} dBm{note}). Data upload starts within a minute.")

    def _go(self, step: Step, command: str, timeout: float) -> None:
        self.step, self._deadline = step, self.now() + timeout
        self._emit(step.value + "...")
        if not self.send(command):
            self._fail("The box is not connected. Connect it first.")

    def _emit(self, text: str, ok: bool | None = None) -> None:
        ev = FlowEvent(self.step, text, ok)
        self._events.append(ev)
        self.on_event(ev)

    def _finish(self, text: str) -> None:
        self.step = Step.DONE
        self._emit(text, True)

    def _fail(self, text: str) -> None:
        self.step = Step.FAILED
        self._emit(text, False)