ninlight-w600/tools/station/tests/test_server.py

import hashlib
import http.client
import json
import socket
import subprocess
import sys
import time
from pathlib import Path

import pytest

SERVER = Path(__file__).resolve().parents[1] / "server.py"
TRAILER_SIZE = 512


def make_image(path: Path, device: str, version: str, payload: bytes) -> bytes:
    """Builds a file with the same trailer layout the firmware verifies."""
    trailer = bytearray(b"\xff" * TRAILER_SIZE)
    trailer[0:4] = b"NINL"
    trailer[4] = 1
    trailer[5:21] = device.encode().ljust(16, b"\x00")
    trailer[21:37] = version.encode().ljust(16, b"\x00")
    trailer[37:69] = hashlib.sha256(payload).digest()
    data = payload + bytes(trailer)
    path.write_bytes(data)
    return data


def free_port() -> int:
    s = socket.socket()
    s.bind(("127.0.0.1", 0))
    port = s.getsockname()[1]
    s.close()
    return port


class Station:
    def __init__(self, tmp_path, *extra):
        self.image_dir = tmp_path / "images"
        self.image_dir.mkdir()
        self.state = tmp_path / "status.json"
        self.port = free_port()
        self.proc = subprocess.Popen(
            [sys.executable, str(SERVER), "--image-dir", str(self.image_dir),
             "--bind", "127.0.0.1", "--port", str(self.port),
             "--state", str(self.state), *extra],
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
        )
        deadline = time.time() + 10
        while time.time() < deadline:
            try:
                socket.create_connection(("127.0.0.1", self.port), 0.2).close()
                return
            except OSError:
                if self.proc.poll() is not None:
                    raise RuntimeError("server died: " + (self.proc.stdout.read() or ""))
                time.sleep(0.05)
        raise RuntimeError("server did not come up")

    def conn(self):
        return http.client.HTTPConnection("127.0.0.1", self.port, timeout=5)

    def read_state(self):
        for _ in range(50):
            if self.state.exists():
                try:
                    return json.loads(self.state.read_text())
                except json.JSONDecodeError:
                    pass
            time.sleep(0.05)
        raise AssertionError("no state file written")

    def stop(self):
        self.proc.terminate()
        self.proc.wait(timeout=5)


@pytest.fixture
def station(tmp_path):
    s = Station(tmp_path)
    yield s
    s.stop()


def test_serves_a_known_image(station):
    data = make_image(station.image_dir / "ninlight_evk_ota.img", "evk", "1.2.3", b"payload" * 500)
    c = station.conn()
    c.request("GET", "/ninlight_evk_ota.img")
    r = c.getresponse()
    body = r.read()
    assert r.status == 200
    assert int(r.getheader("Content-Length")) == len(data)
    assert body == data


def test_missing_image_is_404(station):
    c = station.conn()
    c.request("GET", "/ninlight_g95_cw_ota.img")
    assert c.getresponse().status == 404


def test_unknown_target_is_404(station):
    make_image(station.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"x" * 100)
    for path in ("/ninlight_nosuch_ota.img", "/ninlight_evk_ota.img.bak", "/", "/images/"):
        c = station.conn()
        c.request("GET", path)
        assert c.getresponse().status == 404, path


def test_path_traversal_is_refused(station):
    secret = station.image_dir.parent / "secret.txt"
    secret.write_text("nicht ausliefern")
    c = station.conn()
    c.request("GET", "/../secret.txt")
    assert c.getresponse().status == 404


def test_other_methods_are_405(station):
    c = station.conn()
    c.request("POST", "/ninlight_evk_ota.img", body=b"")
    assert c.getresponse().status == 405


def test_range_request_returns_206_and_the_right_slice(station):
    data = make_image(station.image_dir / "ninlight_e27_rgb_ota.img", "e27_rgb", "2.0.0", b"abc" * 2000)
    c = station.conn()
    c.request("GET", "/ninlight_e27_rgb_ota.img", headers={"Range": "bytes=1000-"})
    r = c.getresponse()
    body = r.read()
    assert r.status == 206
    assert r.getheader("Content-Range") == f"bytes 1000-{len(data)-1}/{len(data)}"
    assert body == data[1000:]


def test_state_file_records_a_delivery(station):
    make_image(station.image_dir / "ninlight_evk_ota.img", "evk", "0.3.4", b"z" * 1000)
    c = station.conn()
    c.request("GET", "/ninlight_evk_ota.img")
    c.getresponse().read()
    state = station.read_state()
    assert state["served"] == 1
    entry = state["recent"][-1]
    assert entry["target"] == "evk"
    assert entry["version"] == "0.3.4"
    assert entry["status"] == "ok"
    assert entry["ip"] == "127.0.0.1"
    assert entry["bytes"] == 1000 + TRAILER_SIZE


def test_state_file_records_a_missing_image(station):
    c = station.conn()
    c.request("GET", "/ninlight_g95_cw_ota.img")
    c.getresponse().read()
    state = station.read_state()
    assert state["served"] == 0
    assert state["recent"][-1]["status"] == "notfound"
    assert state["recent"][-1]["target"] == "g95_cw"


def test_recent_is_capped_at_eight(station):
    make_image(station.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"q" * 10)
    for _ in range(12):
        c = station.conn()
        c.request("GET", "/ninlight_evk_ota.img")
        c.getresponse().read()
    state = station.read_state()
    assert state["served"] == 12
    assert len(state["recent"]) == 8


def test_aborted_transfer_is_recorded(tmp_path):
    s = Station(tmp_path, "--stall-after", "256")
    try:
        make_image(s.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"y" * 100000)
        sock = socket.create_connection(("127.0.0.1", s.port), timeout=5)
        sock.sendall(b"GET /ninlight_evk_ota.img HTTP/1.1\r\nHost: x\r\n\r\n")
        sock.recv(4096)
        time.sleep(0.3)
        sock.close()
        deadline = time.time() + 10
        while time.time() < deadline:
            try:
                entry = json.loads(s.state.read_text())["recent"][-1]
                if entry["status"] == "aborted":
                    assert entry["bytes"] < 100000
                    return
            except (OSError, json.JSONDecodeError, IndexError, KeyError):
                pass
            time.sleep(0.1)
        raise AssertionError("abort was not recorded")
    finally:
        s.stop()


def test_stall_after_stops_sending_but_keeps_the_connection(tmp_path):
    s = Station(tmp_path, "--stall-after", "512")
    try:
        make_image(s.image_dir / "ninlight_evk_ota.img", "evk", "1.0.0", b"y" * 200000)
        sock = socket.create_connection(("127.0.0.1", s.port), timeout=5)
        sock.sendall(b"GET /ninlight_evk_ota.img HTTP/1.1\r\nHost: x\r\n\r\n")
        got = b""
        sock.settimeout(2)
        try:
            while len(got) < 200000:
                chunk = sock.recv(65536)
                if not chunk:
                    break
                got += chunk
        except socket.timeout:
            pass
        # far less than the full image, and the peer never closed the connection
        assert len(got) < 20000
        sock.close()
    finally:
        s.stop()