ninlight-w600/tools/station/server.py

#!/usr/bin/env python3
"""ninlight recovery station.

Serves firmware images to lamps that look for a rescue station while booting. Deliberately
minimal: it hands out nothing but files whose name matches ninlight_<target>_ota.img, so a
stray file in the image directory can never leave the machine.

Standard library only - the Pi gets no extra packages.
"""
from __future__ import annotations

import argparse
import json
import os
import re
import select
import socket
import sys
import tempfile
import threading
import time
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

TARGETS = ("evk", "e27_rgb", "e14_rgb", "gu10_rgb", "e27_cw", "e14_cw", "gu10_cw", "g95_cw")
PATH_RE = re.compile(r"^/ninlight_(" + "|".join(TARGETS) + r")_ota\.img$")
TRAILER_SIZE = 512
TRAILER_MAGIC = b"NINL"
VERSION_OFF = 21
VERSION_LEN = 16
RECENT_MAX = 8
CHUNK = 32 * 1024


def now_iso() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def trailer_version(path: Path) -> str:
    """Firmware version from the image trailer, or '?' if the file carries none."""
    try:
        with path.open("rb") as f:
            if path.stat().st_size < TRAILER_SIZE:
                return "?"
            f.seek(-TRAILER_SIZE, os.SEEK_END)
            tail = f.read(TRAILER_SIZE)
    except OSError:
        return "?"
    if tail[0:4] != TRAILER_MAGIC:
        return "?"
    return tail[VERSION_OFF:VERSION_OFF + VERSION_LEN].rstrip(b"\x00").decode("ascii", "replace")


class Station:
    def __init__(self, image_dir: Path, state_path: Path, stall_after: int | None) -> None:
        self.image_dir = image_dir
        self.state_path = state_path
        self.stall_after = stall_after
        self.since = now_iso()
        self.served = 0
        self.recent: list[dict] = []
        # ThreadingHTTPServer runs one handler thread per connection, and several lamps
        # can be downloading at once. Every mutation of served/recent and the snapshot
        # that gets serialized to disk must happen under this same lock, or concurrent
        # requests lose increments and stomp on each other's status.json writes.
        self.lock = threading.RLock()
        self.write_state()

    def image_path(self, target: str) -> Path:
        return self.image_dir / f"ninlight_{target}_ota.img"

    def record(self, ip: str, target: str, version: str, nbytes: int, status: str) -> None:
        with self.lock:
            if status == "ok":
                self.served += 1
            self.recent.append({"ts": now_iso(), "ip": ip, "target": target,
                                "version": version, "bytes": nbytes, "status": status})
            del self.recent[:-RECENT_MAX]
            self.write_state()

    def write_state(self) -> None:
        with self.lock:
            payload = {"since": self.since, "served": self.served, "recent": list(self.recent)}
            try:
                self.state_path.parent.mkdir(parents=True, exist_ok=True)
                fd, tmp = tempfile.mkstemp(dir=self.state_path.parent, prefix=".status-")
                with os.fdopen(fd, "w") as f:
                    json.dump(payload, f, indent=2)
                os.replace(tmp, self.state_path)
            except OSError as e:
                print(f"warning: cannot write {self.state_path}: {e}", flush=True)


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version = "ninlight-station"

    @property
    def station(self) -> Station:
        return self.server.station  # type: ignore[attr-defined]

    def log_message(self, fmt: str, *args) -> None:  # noqa: D102 - quieter default log
        pass

    def log(self, status: int, nbytes: int, started: float) -> None:
        print(f"{now_iso()} {self.client_address[0]} {self.path} {status} "
              f"{nbytes}B {time.monotonic() - started:.1f}s", flush=True)

    def do_POST(self) -> None:
        self.reject(405, "method not allowed")

    do_PUT = do_DELETE = do_HEAD = do_POST

    def reject(self, status: int, text: str) -> None:
        body = (text + "\n").encode()
        self.send_response(status)
        self.send_header("Content-Type", "text/plain")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Connection", "close")
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self) -> None:
        started = time.monotonic()
        m = PATH_RE.match(self.path)
        if not m:
            self.reject(404, "not found")
            self.log(404, 0, started)
            return
        target = m.group(1)
        path = self.station.image_path(target)
        if not path.is_file():
            # Record before responding: a client that sees the 404 complete may check
            # status.json right away, so the entry must already be on disk by then.
            self.station.record(self.client_address[0], target, "?", 0, "notfound")
            self.reject(404, "no image for this device")
            self.log(404, 0, started)
            return

        size = path.stat().st_size
        start = 0
        rng = self.headers.get("Range", "")
        if rng.startswith("bytes="):
            first = rng.split("=", 1)[1].split("-", 1)[0]
            if first.isdigit() and int(first) < size:
                start = int(first)

        version = trailer_version(path)
        self.send_response(206 if start else 200)
        self.send_header("Content-Type", "application/octet-stream")
        self.send_header("Content-Length", str(size - start))
        if start:
            self.send_header("Content-Range", f"bytes {start}-{size - 1}/{size}")
        self.send_header("Accept-Ranges", "bytes")
        self.send_header("Connection", "close")
        self.end_headers()

        sent = 0
        status = "ok"
        remaining = size - start
        try:
            with path.open("rb") as f:
                f.seek(start)
                while sent < remaining:
                    if self.station.stall_after is not None and sent >= self.station.stall_after:
                        # test aid: hold the connection open and send nothing more, so the
                        # firmware's stall handling can be exercised on real hardware. We
                        # still notice if the peer gives up and closes its side, so a test
                        # (or an impatient client) that hangs up is recorded as "aborted"
                        # instead of leaking the thread forever.
                        while True:
                            readable, _, _ = select.select([self.connection], [], [], 1.0)
                            if readable:
                                try:
                                    peek = self.connection.recv(1, socket.MSG_PEEK)
                                except OSError:
                                    peek = b""
                                if not peek:
                                    raise ConnectionResetError("peer closed during stall")
                    read_size = CHUNK
                    if self.station.stall_after is not None:
                        # Don't overshoot the stall boundary within a single read: without
                        # this a whole CHUNK could be sent before the check above ever
                        # notices we've passed --stall-after.
                        read_size = min(read_size, max(1, self.station.stall_after - sent))
                    chunk = f.read(read_size)
                    if not chunk:
                        break
                    self.wfile.write(chunk)
                    sent += len(chunk)
        except (BrokenPipeError, ConnectionResetError, TimeoutError):
            status = "aborted"
        # Record only now that the transfer has actually finished (or failed): "ok" must
        # never be written before the last byte is confirmed sent, or status.json can
        # claim success for a device that never received the end of the image.
        self.log(206 if start else 200, sent, started)
        self.station.record(self.client_address[0], target, version, sent, status)


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="ninlight recovery station HTTP server")
    p.add_argument("--image-dir", required=True, type=Path)
    p.add_argument("--bind", default="192.168.8.1")
    p.add_argument("--port", type=int, default=80)
    p.add_argument("--state", type=Path, default=Path("/run/ninlight-station/status.json"))
    p.add_argument("--stall-after", type=int, default=None,
                   help="test aid: stop sending after N bytes but keep the connection open")
    a = p.parse_args(argv)

    if not a.image_dir.is_dir():
        print(f"error: image directory {a.image_dir} does not exist", file=sys.stderr)
        return 2

    station = Station(a.image_dir, a.state, a.stall_after)
    httpd = ThreadingHTTPServer((a.bind, a.port), Handler)
    httpd.daemon_threads = True
    httpd.station = station  # type: ignore[attr-defined]
    images = sorted(p.name for p in a.image_dir.glob("ninlight_*_ota.img"))
    print(f"{now_iso()} listening on {a.bind}:{a.port}, {len(images)} image(s): "
          f"{', '.join(images) or 'none'}", flush=True)
    if a.stall_after is not None:
        print(f"{now_iso()} WARNING: --stall-after {a.stall_after} is a test aid", flush=True)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass
    return 0


if __name__ == "__main__":
    raise SystemExit(main())