ninlight-w600/tools/mkimage/test_nl_mkimage.py

import hashlib
import subprocess
import sys
from pathlib import Path

SCRIPT = Path(__file__).with_name("nl_mkimage.py")


def build(tmp_path, payload=b"\x01\x02\x03" * 100, device="evk", version="1.2.3"):
    src = tmp_path / "fw_gz.img"
    src.write_bytes(payload)
    out = tmp_path / "fw_ota.img"
    subprocess.run(
        [sys.executable, str(SCRIPT), "--image", str(src), "--device", device,
         "--version", version, "--out", str(out)],
        check=True,
    )
    return src, out


def test_appends_512_byte_trailer(tmp_path):
    src, out = build(tmp_path)
    data = out.read_bytes()
    assert len(data) == len(src.read_bytes()) + 512
    assert data[: len(src.read_bytes())] == src.read_bytes()
    trailer = data[-512:]
    assert trailer[0:4] == b"NINL"
    assert trailer[4] == 1
    assert trailer[5:21].rstrip(b"\x00") == b"evk"
    assert trailer[21:37].rstrip(b"\x00") == b"1.2.3"
    assert trailer[37:69] == hashlib.sha256(src.read_bytes()).digest()
    assert set(trailer[69:]) == {0xFF}


def test_rejects_long_fields(tmp_path):
    src = tmp_path / "fw_gz.img"
    src.write_bytes(b"x")
    r = subprocess.run(
        [sys.executable, str(SCRIPT), "--image", str(src), "--device", "0123456789abcdef",
         "--version", "1", "--out", str(tmp_path / "o.img")],
        capture_output=True, text=True,
    )
    assert r.returncode != 0
    assert "device" in r.stderr


def test_refuses_image_that_already_has_a_trailer(tmp_path):
    _, out = build(tmp_path)
    r = subprocess.run(
        [sys.executable, str(SCRIPT), "--image", str(out), "--device", "evk",
         "--version", "1.2.3", "--out", str(tmp_path / "again.img")],
        capture_output=True, text=True,
    )
    assert r.returncode != 0
    assert "trailer" in r.stderr


def verify(path, *extra):
    return subprocess.run(
        [sys.executable, str(SCRIPT), "--verify", str(path), *extra],
        capture_output=True, text=True,
    )


def test_verify_accepts_a_good_image(tmp_path):
    _, out = build(tmp_path, device="e27_rgb", version="1.2.3")
    r = verify(out)
    assert r.returncode == 0, r.stderr
    assert r.stdout.split()[:2] == ["e27_rgb", "1.2.3"]


def test_verify_rejects_a_file_without_trailer(tmp_path):
    p = tmp_path / "plain.img"
    p.write_bytes(b"x" * 5000)
    r = verify(p)
    assert r.returncode == 1
    assert "trailer" in r.stderr.lower()


def test_verify_rejects_a_corrupted_payload(tmp_path):
    _, out = build(tmp_path)
    d = bytearray(out.read_bytes())
    d[100] ^= 0x01
    out.write_bytes(d)
    r = verify(out)
    assert r.returncode == 1
    assert "sha256" in r.stderr.lower()


def test_verify_checks_the_expected_device(tmp_path):
    _, out = build(tmp_path, device="evk")
    assert verify(out, "--expect-device", "evk").returncode == 0
    r = verify(out, "--expect-device", "g95_cw")
    assert r.returncode == 1
    assert "evk" in r.stderr and "g95_cw" in r.stderr


def test_verify_rejects_a_file_shorter_than_the_trailer(tmp_path):
    p = tmp_path / "tiny.img"
    p.write_bytes(b"NINL" + b"\x00" * 10)
    r = verify(p)
    assert r.returncode == 1


def test_verify_rejects_nonexistent_file(tmp_path):
    p = tmp_path / "does_not_exist.img"
    r = verify(p)
    assert r.returncode == 1
    assert r.stdout == ""
    assert len(r.stderr.splitlines()) == 1
    assert "error:" in r.stderr.lower()


def test_verify_rejects_directory(tmp_path):
    d = tmp_path / "dir"
    d.mkdir()
    r = verify(d)
    assert r.returncode == 1
    assert r.stdout == ""
    assert len(r.stderr.splitlines()) == 1
    assert "error:" in r.stderr.lower()


def test_verify_rejects_empty_file(tmp_path):
    p = tmp_path / "empty.img"
    p.write_bytes(b"")
    r = verify(p)
    assert r.returncode == 1
    assert "trailer" in r.stderr.lower()


def test_verify_rejects_file_with_only_trailer_no_payload(tmp_path):
    p = tmp_path / "only_trailer.img"
    trailer = bytearray(b"\xff" * 512)
    trailer[0:4] = b"NINL"
    trailer[4] = 1
    trailer[5:21] = b"test\x00" + b"\x00" * 11
    trailer[21:37] = b"1.0.0\x00" + b"\x00" * 10
    trailer[37:69] = hashlib.sha256(b"").digest()
    p.write_bytes(bytes(trailer))
    r = verify(p)
    assert r.returncode == 1
    assert "trailer" in r.stderr.lower()


def test_build_rejects_device_with_whitespace(tmp_path):
    src = tmp_path / "fw_gz.img"
    src.write_bytes(b"x")
    r = subprocess.run(
        [sys.executable, str(SCRIPT), "--image", str(src), "--device", "dev ice",
         "--version", "1.0", "--out", str(tmp_path / "o.img")],
        capture_output=True, text=True,
    )
    assert r.returncode != 0
    assert "whitespace" in r.stderr


def test_build_rejects_version_with_whitespace(tmp_path):
    src = tmp_path / "fw_gz.img"
    src.write_bytes(b"x")
    r = subprocess.run(
        [sys.executable, str(SCRIPT), "--image", str(src), "--device", "device",
         "--version", "1.0 beta", "--out", str(tmp_path / "o.img")],
        capture_output=True, text=True,
    )
    assert r.returncode != 0
    assert "whitespace" in r.stderr


def test_verify_rejects_unterminated_device_field(tmp_path):
    p = tmp_path / "bad_device.img"
    payload = b"x" * 1000
    trailer = bytearray(b"\xff" * 512)
    trailer[0:4] = b"NINL"
    trailer[4] = 1
    # Fill entire device field with no null terminator
    trailer[5:21] = b"a" * 16
    trailer[21:37] = b"1.0.0\x00" + b"\x00" * 10
    trailer[37:69] = hashlib.sha256(payload).digest()
    p.write_bytes(payload + bytes(trailer))
    r = verify(p)
    assert r.returncode == 1
    assert "unterminated" in r.stderr.lower()


def test_verify_rejects_unterminated_device_field_without_expect_device(tmp_path):
    # Same as above but without --expect-device to ensure the check happens regardless
    p = tmp_path / "bad_device2.img"
    payload = b"x" * 1000
    trailer = bytearray(b"\xff" * 512)
    trailer[0:4] = b"NINL"
    trailer[4] = 1
    # Fill entire device field with no null terminator
    trailer[5:21] = b"b" * 16
    trailer[21:37] = b"2.0.0\x00" + b"\x00" * 10
    trailer[37:69] = hashlib.sha256(payload).digest()
    p.write_bytes(payload + bytes(trailer))
    r = verify(p)
    assert r.returncode == 1
    assert "unterminated" in r.stderr.lower()


def test_verify_with_build_arguments_is_an_error(tmp_path):
    _, out = build(tmp_path)
    r = subprocess.run(
        [sys.executable, str(SCRIPT), "--verify", str(out), "--device", "x"],
        capture_output=True, text=True,
    )
    assert r.returncode != 0
    assert "cannot be combined" in r.stderr


# --- Python 3.8 compatibility -------------------------------------------------------------
#
# CMake calls this script via plain `python3` inside the build container on every firmware
# build (see docker/Dockerfile, cmake/run_mkimage.cmake), and that container is stuck on
# Python 3.8.10. This is the opposite constraint from the rest of the repo's tooling (which
# needs Python >= 3.11) - see the module docstring in nl_mkimage.py for the full story, and
# git history around the "unsupported operand type(s) for |" break for what happens when it
# regresses.
#
# Two checks, because neither alone covers the failure mode we actually hit:
#   1. `from __future__ import annotations` must be the first statement after the module
#      docstring. As long as it's there, PEP 604 unions in annotations (`str | None`) are
#      never evaluated at runtime, so they can't raise on 3.8/3.9 no matter how many more of
#      them get added later.
#   2. vermin's *default* analysis (annotations not evaluated - matching point 1) catches
#      everything else 3.9+/3.10+/3.11+ that could sneak in: walrus in odd positions, match
#      statements, dict `|=`, new stdlib members, etc. We deliberately do NOT enable vermin's
#      `union-types` feature: it flags `str | None` as needing 3.10 unconditionally, without
#      accounting for `from __future__ import annotations` deferring evaluation - which would
#      make this test fail on the very fix that makes the file 3.8-safe. Check 1 above is
#      what actually guards against that pattern.


def test_future_annotations_import_is_present():
    """Guards the PEP 604 union pattern that broke the Docker build (see module docstring)."""
    import ast

    tree = ast.parse(SCRIPT.read_text())
    body = tree.body
    # Skip a leading module docstring, if any.
    if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant) \
            and isinstance(body[0].value.value, str):
        body = body[1:]
    assert body, "nl_mkimage.py is unexpectedly empty"
    first = body[0]
    assert isinstance(first, ast.ImportFrom) and first.module == "__future__" \
        and any(alias.name == "annotations" for alias in first.names), (
        "nl_mkimage.py must start with `from __future__ import annotations` right after its "
        "module docstring - without it, PEP 604 unions like `str | None` in annotations are "
        "evaluated eagerly and crash under the Python 3.8.10 in the Docker build container "
        "(TypeError: unsupported operand type(s) for |: 'type' and 'NoneType')"
    )


def test_stays_python38_compatible():
    """Static check that nl_mkimage.py needs nothing newer than Python 3.8 to run.

    Runs without a Python 3.8 interpreter installed: vermin infers the minimum required
    version from syntax/stdlib usage instead of actually executing the file.
    """
    vermin = __import__("vermin")

    config = vermin.Config()
    config.set_quiet(True)
    # Match nl_mkimage.py's own strategy: with `from __future__ import annotations` present
    # (checked separately above), annotations are never evaluated, so vermin shouldn't assume
    # they need runtime support either. This is vermin's default, spelled out here so it's
    # obvious why `union-types` (see the block comment above) is intentionally not enabled.
    config.set_eval_annotations(False)

    source = SCRIPT.read_text()
    py2_min, py3_min = vermin.detect(source, config, str(SCRIPT))
    assert py3_min is not None and py3_min <= (3, 8), (
        f"nl_mkimage.py now needs Python {py3_min} - it must keep working on the "
        "Python 3.8.10 in the Docker build container (see module docstring)"
    )