U1 bootlogo injector

inject_boot_logo.py

#!/usr/bin/env python3
"""Safely replace the two Rockchip boot logos in a Snapmaker U1 boot slot.

The script uses only the Python standard library. By default it performs a
dry-run: the source boot image or selected boot partition is read, validated,
patched in memory, and written to a regular staging file. A block device is
written only when --apply is supplied.
"""

from __future__ import annotations

import argparse
import fcntl
import hashlib
import os
import shutil
import stat
import struct
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable


EXPECTED_BMP_WIDTH = 480
EXPECTED_BMP_HEIGHT = 320
EXPECTED_BMP_BPP = 24
EXPECTED_BMP_DIB_SIZE = 40
EXPECTED_BMP_PIXEL_OFFSET = 54
EXPECTED_BMP_PIXEL_BYTES = EXPECTED_BMP_WIDTH * EXPECTED_BMP_HEIGHT * 3
EXPECTED_BMP_FILE_SIZE = EXPECTED_BMP_PIXEL_OFFSET + EXPECTED_BMP_PIXEL_BYTES
EXPECTED_BOOT_PARTITION_SIZE = 32 * 1024 * 1024

FDT_MAGIC = 0xD00DFEED
FDT_BEGIN_NODE = 1
FDT_END_NODE = 2
FDT_PROP = 3
FDT_NOP = 4
FDT_END = 9

RESOURCE_BLOCK_SIZE = 512
RESOURCE_HEADER_MAGIC = b"RSCE"
RESOURCE_ENTRY_MAGIC = b"ENTR"
RESOURCE_ENTRY_PATH_OFFSET = 4
RESOURCE_ENTRY_PATH_SIZE = 220
RESOURCE_ENTRY_HASH_OFFSET = 224
RESOURCE_ENTRY_HASH_CAPACITY = 32
RESOURCE_ENTRY_HASH_SIZE_OFFSET = 256
RESOURCE_ENTRY_CONTENT_OFFSET_OFFSET = 260
RESOURCE_ENTRY_CONTENT_SIZE_OFFSET = 264

BLKGETSIZE64 = 0x80081272


class InjectionError(RuntimeError):
    """A validation or safety check failed."""


@dataclass(frozen=True)
class FdtProperty:
    name: str
    value: bytes
    value_offset: int


@dataclass
class FdtNode:
    path: str
    properties: dict[str, FdtProperty]


@dataclass(frozen=True)
class ResourceEntry:
    path: str
    entry_offset: int
    hash_size: int
    content_block: int
    content_size: int

    @property
    def content_offset(self) -> int:
        return self.content_block * RESOURCE_BLOCK_SIZE


@dataclass(frozen=True)
class Source:
    path: Path
    display_name: str
    slot: str | None
    is_block_device: bool
    size: int


def fail(message: str) -> None:
    raise InjectionError(message)


def align4(value: int) -> int:
    return (value + 3) & ~3


def unpack_u32_be(data: bytes | bytearray, offset: int, description: str) -> int:
    if offset < 0 or offset + 4 > len(data):
        fail(f"Truncated {description}")
    return struct.unpack_from(">I", data, offset)[0]


def decode_c_string(data: bytes, offset: int, description: str) -> str:
    if offset < 0 or offset >= len(data):
        fail(f"Invalid string offset for {description}")
    end = data.find(b"\0", offset)
    if end < 0:
        fail(f"Unterminated string for {description}")
    try:
        return data[offset:end].decode("ascii")
    except UnicodeDecodeError as exc:
        raise InjectionError(f"Non-ASCII string for {description}") from exc


def validate_bmp(path: Path) -> bytes:
    try:
        data = path.read_bytes()
    except OSError as exc:
        raise InjectionError(f"Cannot read BMP {path}: {exc}") from exc

    if len(data) < EXPECTED_BMP_PIXEL_OFFSET:
        fail(f"BMP is truncated: {len(data)} bytes")
    if data[:2] != b"BM":
        fail("BMP signature is not 'BM'")

    declared_size = struct.unpack_from("<I", data, 2)[0]
    pixel_offset = struct.unpack_from("<I", data, 10)[0]
    dib_size = struct.unpack_from("<I", data, 14)[0]
    width = struct.unpack_from("<i", data, 18)[0]
    height = struct.unpack_from("<i", data, 22)[0]
    planes = struct.unpack_from("<H", data, 26)[0]
    bits_per_pixel = struct.unpack_from("<H", data, 28)[0]
    compression = struct.unpack_from("<I", data, 30)[0]
    declared_pixel_size = struct.unpack_from("<I", data, 34)[0]

    if dib_size != EXPECTED_BMP_DIB_SIZE:
        fail(
            f"BMP DIB header is {dib_size} bytes; expected "
            f"{EXPECTED_BMP_DIB_SIZE}"
        )
    if pixel_offset != EXPECTED_BMP_PIXEL_OFFSET:
        fail(
            f"BMP pixel offset is {pixel_offset}; expected "
            f"{EXPECTED_BMP_PIXEL_OFFSET}"
        )
    if width != EXPECTED_BMP_WIDTH or abs(height) != EXPECTED_BMP_HEIGHT:
        fail(
            f"BMP resolution is {width}x{abs(height)}; expected "
            f"{EXPECTED_BMP_WIDTH}x{EXPECTED_BMP_HEIGHT}"
        )
    if planes != 1:
        fail(f"BMP has {planes} color planes; expected 1")
    if bits_per_pixel != EXPECTED_BMP_BPP:
        fail(
            f"BMP is {bits_per_pixel}-bit; expected "
            f"{EXPECTED_BMP_BPP}-bit"
        )
    if compression != 0:
        fail(f"BMP compression mode is {compression}; expected uncompressed")
    if declared_pixel_size not in (0, EXPECTED_BMP_PIXEL_BYTES):
        fail(
            f"BMP declares {declared_pixel_size} pixel bytes; expected "
            f"{EXPECTED_BMP_PIXEL_BYTES}"
        )
    if declared_size != len(data):
        fail(
            f"BMP header declares {declared_size} bytes, but the file is "
            f"{len(data)} bytes"
        )
    if len(data) != EXPECTED_BMP_FILE_SIZE:
        fail(
            f"BMP file is {len(data)} bytes; expected exactly "
            f"{EXPECTED_BMP_FILE_SIZE} bytes"
        )

    orientation = "top-down" if height < 0 else "bottom-up"
    print(
        f"BMP validated: {width}x{abs(height)}, {bits_per_pixel}-bit, "
        f"uncompressed, {orientation}, {len(data)} bytes"
    )
    return data


def parse_fdt(data: bytes | bytearray) -> dict[str, FdtNode]:
    if len(data) < 40:
        fail("Boot image is too small to contain an FDT header")

    (
        magic,
        total_size,
        struct_offset,
        strings_offset,
        _reserve_map_offset,
        version,
        last_compatible_version,
        _boot_cpu_id,
        strings_size,
        struct_size,
    ) = struct.unpack_from(">10I", data, 0)

    if magic != FDT_MAGIC:
        fail(f"Invalid FIT/FDT magic 0x{magic:08x}")
    if version < 17 or last_compatible_version > 17:
        fail(
            f"Unsupported FDT version {version} "
            f"(last compatible {last_compatible_version})"
        )
    if total_size > len(data):
        fail(f"FDT total size {total_size} exceeds source size {len(data)}")
    if struct_offset + struct_size > total_size:
        fail("FDT structure block is outside the FDT")
    if strings_offset + strings_size > total_size:
        fail("FDT strings block is outside the FDT")

    strings = bytes(data[strings_offset : strings_offset + strings_size])
    nodes: dict[str, FdtNode] = {}
    stack: list[str] = []
    position = struct_offset
    structure_end = struct_offset + struct_size
    saw_end = False

    while position < structure_end:
        token = unpack_u32_be(data, position, "FDT token")
        position += 4

        if token == FDT_BEGIN_NODE:
            name_end = bytes(data).find(b"\0", position, structure_end)
            if name_end < 0:
                fail("Unterminated FDT node name")
            try:
                name = bytes(data[position:name_end]).decode("ascii")
            except UnicodeDecodeError as exc:
                raise InjectionError("Non-ASCII FDT node name") from exc
            position = align4(name_end + 1)

            if not stack:
                if name:
                    fail("FDT root node has a non-empty name")
                path = "/"
            else:
                path = stack[-1].rstrip("/") + "/" + name
            if path in nodes:
                fail(f"Duplicate FDT node {path}")
            nodes[path] = FdtNode(path=path, properties={})
            stack.append(path)

        elif token == FDT_END_NODE:
            if not stack:
                fail("Unexpected FDT_END_NODE")
            stack.pop()

        elif token == FDT_PROP:
            if not stack:
                fail("FDT property outside a node")
            length = unpack_u32_be(data, position, "FDT property length")
            name_offset = unpack_u32_be(
                data, position + 4, "FDT property name offset"
            )
            position += 8
            if position + length > structure_end:
                fail("FDT property extends outside the structure block")
            name = decode_c_string(strings, name_offset, "FDT property")
            node = nodes[stack[-1]]
            if name in node.properties:
                fail(f"Duplicate property {name} in {node.path}")
            node.properties[name] = FdtProperty(
                name=name,
                value=bytes(data[position : position + length]),
                value_offset=position,
            )
            position = align4(position + length)

        elif token == FDT_NOP:
            continue

        elif token == FDT_END:
            if stack:
                fail("FDT ended before all nodes were closed")
            saw_end = True
            break

        else:
            fail(f"Unknown FDT token 0x{token:08x}")

    if not saw_end:
        fail("FDT structure has no FDT_END token")
    return nodes


def immediate_children(
    nodes: dict[str, FdtNode], parent_path: str
) -> list[FdtNode]:
    prefix = parent_path.rstrip("/") + "/"
    return [
        node
        for path, node in nodes.items()
        if path.startswith(prefix) and "/" not in path[len(prefix) :]
    ]


def required_property(node: FdtNode, name: str) -> FdtProperty:
    prop = node.properties.get(name)
    if prop is None:
        fail(f"Required FIT property {node.path}:{name} is missing")
    return prop


def property_u32(node: FdtNode, name: str) -> int:
    prop = required_property(node, name)
    if len(prop.value) != 4:
        fail(f"FIT property {node.path}:{name} is not a 32-bit value")
    return struct.unpack(">I", prop.value)[0]


def property_string(node: FdtNode, name: str) -> str:
    prop = required_property(node, name)
    if not prop.value.endswith(b"\0"):
        fail(f"FIT property {node.path}:{name} is not NUL-terminated")
    try:
        return prop.value[:-1].decode("ascii")
    except UnicodeDecodeError as exc:
        raise InjectionError(
            f"FIT property {node.path}:{name} is not ASCII"
        ) from exc


def find_sha256_hash_node(
    nodes: dict[str, FdtNode], image_node: FdtNode
) -> FdtNode:
    matching: list[FdtNode] = []
    for child in immediate_children(nodes, image_node.path):
        if not child.path.rsplit("/", 1)[-1].startswith("hash"):
            continue
        if property_string(child, "algo") == "sha256":
            matching.append(child)
    if len(matching) != 1:
        fail(
            f"Expected exactly one SHA-256 hash below {image_node.path}; "
            f"found {len(matching)}"
        )
    value = required_property(matching[0], "value")
    if len(value.value) != hashlib.sha256().digest_size:
        fail(f"SHA-256 value below {image_node.path} is not 32 bytes")
    return matching[0]


def validate_fit(
    data: bytes | bytearray,
) -> tuple[dict[str, FdtNode], FdtNode, FdtProperty, int, int]:
    nodes = parse_fdt(data)
    images = nodes.get("/images")
    configurations = nodes.get("/configurations")
    if images is None or configurations is None:
        fail("FIT does not contain /images and /configurations")

    image_nodes = immediate_children(nodes, "/images")
    image_names = {node.path.rsplit("/", 1)[-1] for node in image_nodes}
    if image_names != {"fdt", "kernel", "resource"}:
        fail(
            "Unexpected FIT image set: "
            + ", ".join(sorted(image_names))
            + " (expected fdt, kernel, resource)"
        )

    for path, node in nodes.items():
        if not path.startswith("/configurations/"):
            continue
        basename = path.rsplit("/", 1)[-1]
        if basename.startswith("signature"):
            fail(
                f"Signed FIT configuration detected at {path}; refusing to "
                "invalidate a signature"
            )
        algo = node.properties.get("algo")
        if algo and b"rsa" in algo.value.lower():
            fail(
                f"Signed FIT configuration algorithm detected at {path}; "
                "refusing to continue"
            )

    resource_node = nodes["/images/resource"]
    if property_string(resource_node, "type") != "multi":
        fail("FIT resource image type is not 'multi'")
    if property_string(resource_node, "compression") != "none":
        fail("FIT resource image is compressed")

    occupied_ranges: list[tuple[int, int, str]] = []
    resource_hash_node: FdtNode | None = None
    resource_hash_property: FdtProperty | None = None
    resource_position = 0
    resource_size = 0

    for image_node in image_nodes:
        name = image_node.path.rsplit("/", 1)[-1]
        position = property_u32(image_node, "data-position")
        size = property_u32(image_node, "data-size")
        end = position + size
        if position % RESOURCE_BLOCK_SIZE:
            fail(f"FIT image {name} position is not 512-byte aligned")
        if size <= 0 or end > len(data):
            fail(
                f"FIT image {name} range 0x{position:x}-0x{end:x} is "
                "outside the source"
            )
        for other_start, other_end, other_name in occupied_ranges:
            if position < other_end and other_start < end:
                fail(f"FIT images {name} and {other_name} overlap")
        occupied_ranges.append((position, end, name))

        hash_node = find_sha256_hash_node(nodes, image_node)
        hash_property = required_property(hash_node, "value")
        actual_hash = hashlib.sha256(data[position:end]).digest()
        if actual_hash != hash_property.value:
            fail(f"Stored SHA-256 does not match FIT image {name}")

        if name == "resource":
            resource_hash_node = hash_node
            resource_hash_property = hash_property
            resource_position = position
            resource_size = size

    if resource_hash_node is None or resource_hash_property is None:
        fail("FIT resource SHA-256 node was not found")

    return (
        nodes,
        resource_node,
        resource_hash_property,
        resource_position,
        resource_size,
    )


def parse_resource(data: bytes | bytearray) -> list[ResourceEntry]:
    if len(data) < RESOURCE_BLOCK_SIZE:
        fail("Rockchip resource image is truncated")
    if data[:4] != RESOURCE_HEADER_MAGIC:
        fail("Rockchip resource header magic is not RSCE")

    resource_version, table_version = struct.unpack_from("<HH", data, 4)
    header_blocks = data[8]
    table_offset_blocks = data[9]
    entry_blocks = data[10]
    entry_count = struct.unpack_from("<I", data, 12)[0]

    if (resource_version, table_version) != (0, 0):
        fail(
            f"Unsupported resource/table version "
            f"{resource_version}.{table_version}"
        )
    if header_blocks != 1 or table_offset_blocks != 1 or entry_blocks != 1:
        fail(
            "Unsupported Rockchip resource header/table block layout "
            f"({header_blocks}, {table_offset_blocks}, {entry_blocks})"
        )
    if entry_count <= 0 or entry_count > 64:
        fail(f"Unreasonable Rockchip resource entry count {entry_count}")

    entries: list[ResourceEntry] = []
    seen_paths: set[str] = set()
    table_start = table_offset_blocks * RESOURCE_BLOCK_SIZE

    for index in range(entry_count):
        entry_offset = table_start + index * entry_blocks * RESOURCE_BLOCK_SIZE
        if entry_offset + RESOURCE_BLOCK_SIZE > len(data):
            fail(f"Resource index entry {index} is truncated")
        if data[entry_offset : entry_offset + 4] != RESOURCE_ENTRY_MAGIC:
            fail(f"Resource index entry {index} magic is not ENTR")

        path_bytes = bytes(
            data[
                entry_offset
                + RESOURCE_ENTRY_PATH_OFFSET : entry_offset
                + RESOURCE_ENTRY_PATH_OFFSET
                + RESOURCE_ENTRY_PATH_SIZE
            ]
        )
        nul = path_bytes.find(b"\0")
        if nul < 0:
            fail(f"Resource index entry {index} path is not terminated")
        try:
            path = path_bytes[:nul].decode("ascii")
        except UnicodeDecodeError as exc:
            raise InjectionError(
                f"Resource index entry {index} path is not ASCII"
            ) from exc
        if not path or path in seen_paths:
            fail(f"Invalid or duplicate resource path {path!r}")
        seen_paths.add(path)

        hash_size = struct.unpack_from(
            "<I", data, entry_offset + RESOURCE_ENTRY_HASH_SIZE_OFFSET
        )[0]
        content_block = struct.unpack_from(
            "<I", data, entry_offset + RESOURCE_ENTRY_CONTENT_OFFSET_OFFSET
        )[0]
        content_size = struct.unpack_from(
            "<I", data, entry_offset + RESOURCE_ENTRY_CONTENT_SIZE_OFFSET
        )[0]
        if hash_size not in (20, 32):
            fail(f"Resource {path} has unsupported hash size {hash_size}")
        content_offset = content_block * RESOURCE_BLOCK_SIZE
        if (
            content_size <= 0
            or content_offset < table_start + entry_count * RESOURCE_BLOCK_SIZE
            or content_offset + content_size > len(data)
        ):
            fail(f"Resource {path} has an invalid content range")

        entry = ResourceEntry(
            path=path,
            entry_offset=entry_offset,
            hash_size=hash_size,
            content_block=content_block,
            content_size=content_size,
        )
        content = data[content_offset : content_offset + content_size]
        hash_function = hashlib.sha1 if hash_size == 20 else hashlib.sha256
        actual_hash = hash_function(content).digest()
        stored_hash = bytes(
            data[
                entry_offset
                + RESOURCE_ENTRY_HASH_OFFSET : entry_offset
                + RESOURCE_ENTRY_HASH_OFFSET
                + hash_size
            ]
        )
        if actual_hash != stored_hash:
            fail(f"Stored hash does not match resource {path}")
        entries.append(entry)

    ranges = sorted(
        (
            entry.content_offset,
            entry.content_offset + entry.content_size,
            entry.path,
        )
        for entry in entries
    )
    for index in range(1, len(ranges)):
        if ranges[index][0] < ranges[index - 1][1]:
            fail(
                f"Resource entries {ranges[index - 1][2]} and "
                f"{ranges[index][2]} overlap"
            )

    return entries


def entry_capacity(
    entry: ResourceEntry, entries: Iterable[ResourceEntry], resource_size: int
) -> int:
    later_offsets = [
        candidate.content_offset
        for candidate in entries
        if candidate.content_offset > entry.content_offset
    ]
    end = min(later_offsets) if later_offsets else resource_size
    return end - entry.content_offset


def patch_resource(resource_data: bytes, bmp_data: bytes) -> bytes:
    resource = bytearray(resource_data)
    entries = parse_resource(resource)
    by_path = {entry.path: entry for entry in entries}

    if set(by_path) != {"rk-kernel.dtb", "logo.bmp", "logo_kernel.bmp"}:
        fail(
            "Unexpected resource set: "
            + ", ".join(sorted(by_path))
            + " (expected rk-kernel.dtb, logo.bmp, logo_kernel.bmp)"
        )

    for target in ("logo.bmp", "logo_kernel.bmp"):
        entry = by_path[target]
        if entry.hash_size != hashlib.sha1().digest_size:
            fail(f"Resource {target} does not use the expected SHA-1 hash")
        capacity = entry_capacity(entry, entries, len(resource))
        if len(bmp_data) > capacity:
            fail(
                f"BMP is {len(bmp_data)} bytes but {target} has only "
                f"{capacity} bytes of allocated space"
            )

        clear_size = max(entry.content_size, len(bmp_data))
        start = entry.content_offset
        resource[start : start + clear_size] = b"\0" * clear_size
        resource[start : start + len(bmp_data)] = bmp_data

        digest = hashlib.sha1(bmp_data).digest()
        hash_offset = entry.entry_offset + RESOURCE_ENTRY_HASH_OFFSET
        resource[hash_offset : hash_offset + len(digest)] = digest
        struct.pack_into(
            "<I",
            resource,
            entry.entry_offset + RESOURCE_ENTRY_CONTENT_SIZE_OFFSET,
            len(bmp_data),
        )

    verified_entries = parse_resource(resource)
    verified_by_path = {entry.path: entry for entry in verified_entries}
    for target in ("logo.bmp", "logo_kernel.bmp"):
        entry = verified_by_path[target]
        embedded = bytes(
            resource[
                entry.content_offset : entry.content_offset + entry.content_size
            ]
        )
        if embedded != bmp_data:
            fail(f"Post-patch verification failed for {target}")

    return bytes(resource)


def patch_boot_image(source_data: bytes, bmp_data: bytes) -> bytes:
    (
        _nodes,
        _resource_node,
        resource_hash_property,
        resource_position,
        resource_size,
    ) = validate_fit(source_data)

    resource_end = resource_position + resource_size
    original_resource = source_data[resource_position:resource_end]
    patched_resource = patch_resource(original_resource, bmp_data)
    if len(patched_resource) != resource_size:
        fail("Internal error: resource image size changed")

    patched = bytearray(source_data)
    patched[resource_position:resource_end] = patched_resource
    patched_hash = hashlib.sha256(patched_resource).digest()
    hash_offset = resource_hash_property.value_offset
    patched[hash_offset : hash_offset + len(patched_hash)] = patched_hash

    validate_fit(patched)
    return bytes(patched)


def parse_running_slot(cmdline_path: Path = Path("/proc/cmdline")) -> str:
    try:
        tokens = cmdline_path.read_text(encoding="ascii").split()
    except OSError as exc:
        raise InjectionError(f"Cannot read {cmdline_path}: {exc}") from exc

    accepted_keys = {"android_slotsufix", "androidboot.slot_suffix"}
    discovered: list[tuple[str, str]] = []
    for token in tokens:
        if "=" not in token:
            continue
        key, value = token.split("=", 1)
        if key in accepted_keys:
            discovered.append((key, value))

    if not discovered:
        fail(
            "No android_slotsufix or androidboot.slot_suffix was found in "
            f"{cmdline_path}"
        )
    values = {value for _key, value in discovered}
    if len(values) != 1:
        fail(f"Conflicting slot values in {cmdline_path}: {discovered}")
    value = values.pop()
    if value not in {"_a", "_b"}:
        fail(f"Unsupported active slot value {value!r}")
    return value[1]


def resolve_slot(requested: str) -> tuple[str, str]:
    active = parse_running_slot()
    if requested == "active":
        selected = active
    elif requested == "inactive":
        selected = "b" if active == "a" else "a"
    else:
        selected = requested
    return active, selected


def block_device_size(path: Path) -> int:
    resolved = Path(os.path.realpath(path))
    sysfs_size = Path("/sys/class/block") / resolved.name / "size"
    try:
        sectors = int(sysfs_size.read_text(encoding="ascii").strip())
        size = sectors * 512
        if size > 0:
            return size
    except (OSError, ValueError):
        pass

    try:
        with path.open("rb", buffering=0) as handle:
            buffer = bytearray(8)
            fcntl.ioctl(handle.fileno(), BLKGETSIZE64, buffer, True)
            size = struct.unpack("=Q", buffer)[0]
    except OSError as exc:
        raise InjectionError(
            f"Cannot determine block-device size for {path}: {exc}"
        ) from exc
    if size <= 0:
        fail(f"Block device {path} reports an invalid size {size}")
    return size


def source_from_args(args: argparse.Namespace) -> Source:
    if args.input is not None:
        path = args.input.resolve()
        try:
            mode = path.stat().st_mode
        except OSError as exc:
            raise InjectionError(f"Cannot stat input {path}: {exc}") from exc
        if not stat.S_ISREG(mode):
            fail("--input must name a regular offline image file")
        size = path.stat().st_size
        return Source(
            path=path,
            display_name=str(path),
            slot=None,
            is_block_device=False,
            size=size,
        )

    active, selected = resolve_slot(args.slot)
    path = Path(f"/dev/block/by-name/boot_{selected}")
    try:
        mode = path.stat().st_mode
    except OSError as exc:
        raise InjectionError(f"Cannot stat boot device {path}: {exc}") from exc
    if not stat.S_ISBLK(mode):
        fail(f"Resolved boot source is not a block device: {path}")
    size = block_device_size(path)
    if size != EXPECTED_BOOT_PARTITION_SIZE:
        fail(
            f"Boot device {path} is {size} bytes; expected exactly "
            f"{EXPECTED_BOOT_PARTITION_SIZE} bytes"
        )
    role = "active" if selected == active else "inactive"
    print(
        f"Running slot: {active.upper()}; selected slot: "
        f"{selected.upper()} ({role})"
    )
    return Source(
        path=path,
        display_name=f"{path} ({os.path.realpath(path)})",
        slot=selected,
        is_block_device=True,
        size=size,
    )


def read_exact(path: Path, size: int) -> bytes:
    chunks: list[bytes] = []
    remaining = size
    try:
        with path.open("rb", buffering=0) as handle:
            while remaining:
                chunk = handle.read(min(1024 * 1024, remaining))
                if not chunk:
                    fail(
                        f"Unexpected end of {path}: "
                        f"{remaining} bytes were still expected"
                    )
                chunks.append(chunk)
                remaining -= len(chunk)
    except OSError as exc:
        raise InjectionError(f"Cannot read {path}: {exc}") from exc
    return b"".join(chunks)


def ensure_output_is_safe(output: Path, source: Source, force: bool) -> Path:
    output = output.resolve()
    if output == source.path.resolve():
        fail("Output path must not be the source image or block device")
    if output.exists():
        mode = output.stat().st_mode
        if not stat.S_ISREG(mode):
            fail(f"Output exists and is not a regular file: {output}")
        if not force:
            fail(f"Output already exists: {output} (use --force-output)")
    return output


def write_regular_file(path: Path, data: bytes, force: bool = False) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}")
    if temporary.exists():
        fail(f"Temporary output already exists: {temporary}")
    try:
        flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
        fd = os.open(temporary, flags, 0o600)
        try:
            view = memoryview(data)
            while view:
                written = os.write(fd, view)
                if written <= 0:
                    fail(f"Failed while writing {temporary}")
                view = view[written:]
            os.fsync(fd)
        finally:
            os.close(fd)

        if path.exists() and not force:
            fail(f"Output appeared while staging: {path}")
        os.replace(temporary, path)
        directory_fd = os.open(path.parent, os.O_RDONLY)
        try:
            os.fsync(directory_fd)
        finally:
            os.close(directory_fd)
    except Exception:
        try:
            temporary.unlink()
        except OSError:
            pass
        raise


def default_output(source: Source) -> Path:
    if source.is_block_device:
        assert source.slot is not None
        return Path(f"/tmp/boot_{source.slot}_logo_patched.img")
    return source.path.with_name(source.path.name + ".logo-patched.img")


def check_not_mounted(source: Source) -> None:
    if not source.is_block_device:
        return
    resolved_source = os.path.realpath(source.path)
    try:
        lines = Path("/proc/mounts").read_text(
            encoding="utf-8", errors="replace"
        ).splitlines()
    except OSError as exc:
        raise InjectionError(f"Cannot inspect /proc/mounts: {exc}") from exc
    for line in lines:
        fields = line.split()
        if fields and os.path.realpath(fields[0]) == resolved_source:
            fail(f"Refusing to write mounted block device {source.path}")


def apply_to_device(
    source: Source,
    original: bytes,
    patched: bytes,
    backup_dir: Path,
) -> Path:
    if not source.is_block_device:
        fail("--apply requires an on-device boot-slot source, not --input")
    if os.geteuid() != 0:
        fail("--apply must be run as root")
    check_not_mounted(source)

    backup_dir = backup_dir.resolve()
    backup_dir.mkdir(parents=True, exist_ok=True)
    available = shutil.disk_usage(backup_dir).free
    if available < len(original) + 16 * 1024 * 1024:
        fail(
            f"Insufficient backup space in {backup_dir}: "
            f"{available} bytes free"
        )

    original_hash = hashlib.sha256(original).hexdigest()
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    assert source.slot is not None
    backup = backup_dir / (
        f"boot_{source.slot}_{timestamp}_{original_hash[:12]}.img"
    )
    if backup.exists():
        fail(f"Backup path already exists: {backup}")
    write_regular_file(backup, original)
    backup_data = read_exact(backup, len(original))
    if hashlib.sha256(backup_data).digest() != hashlib.sha256(original).digest():
        fail(f"Backup verification failed: {backup}")
    print(f"Verified backup: {backup}")

    current = read_exact(source.path, source.size)
    if hashlib.sha256(current).digest() != hashlib.sha256(original).digest():
        fail("Boot partition changed after staging; refusing to write")

    print(f"Writing {len(patched)} bytes to {source.display_name}...")
    try:
        fd = os.open(source.path, os.O_WRONLY | os.O_SYNC)
        try:
            os.lseek(fd, 0, os.SEEK_SET)
            view = memoryview(patched)
            while view:
                written = os.write(fd, view)
                if written <= 0:
                    fail(f"Write failed for {source.path}")
                view = view[written:]
            os.fsync(fd)
        finally:
            os.close(fd)
    except OSError as exc:
        raise InjectionError(
            f"Failed to write boot device {source.path}: {exc}. "
            f"Backup is at {backup}"
        ) from exc

    read_back = read_exact(source.path, source.size)
    if hashlib.sha256(read_back).digest() != hashlib.sha256(patched).digest():
        fail(
            f"Read-back verification failed for {source.path}. "
            f"Backup is at {backup}"
        )
    validate_fit(read_back)
    print(f"Verified flashed boot partition: {source.path}")
    return backup


def build_argument_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Replace logo.bmp and logo_kernel.bmp in a Snapmaker U1 boot "
            "image. Dry-run staging is the default; --apply is required to "
            "write a boot partition."
        )
    )
    parser.add_argument("bmp", type=Path, help="custom 480x320 24-bit BMP")
    parser.add_argument(
        "--input",
        type=Path,
        help="offline regular boot image instead of an on-device boot slot",
    )
    parser.add_argument(
        "--slot",
        choices=("active", "inactive", "a", "b"),
        default="active",
        help="on-device slot selection (default: active)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help=(
            "complete staged boot image path; defaults to /tmp on-device"
        ),
    )
    parser.add_argument(
        "--force-output",
        action="store_true",
        help="replace an existing regular staging output",
    )
    mode = parser.add_mutually_exclusive_group()
    mode.add_argument(
        "--dry-run",
        action="store_true",
        help="stage and verify only (this is the default)",
    )
    mode.add_argument(
        "--apply",
        action="store_true",
        help="back up and write the selected on-device boot partition",
    )
    parser.add_argument(
        "--backup-dir",
        type=Path,
        default=Path("/userdata/boot-logo-backups"),
        help="backup directory used with --apply",
    )
    return parser


def main() -> int:
    args = build_argument_parser().parse_args()
    if args.apply and args.input is not None:
        fail("--apply cannot be combined with --input")

    bmp_data = validate_bmp(args.bmp.resolve())
    source = source_from_args(args)
    print(f"Reading source: {source.display_name} ({source.size} bytes)")
    original = read_exact(source.path, source.size)
    original_hash = hashlib.sha256(original).hexdigest()
    print(f"Source SHA-256: {original_hash}")

    patched = patch_boot_image(original, bmp_data)
    if len(patched) != len(original):
        fail("Internal error: complete boot image size changed")
    patched_hash = hashlib.sha256(patched).hexdigest()
    print(f"Patched SHA-256: {patched_hash}")

    output = args.output if args.output is not None else default_output(source)
    output = ensure_output_is_safe(output, source, args.force_output)
    write_regular_file(output, patched, force=args.force_output)
    staged = read_exact(output, len(patched))
    if hashlib.sha256(staged).digest() != hashlib.sha256(patched).digest():
        fail(f"Staged output verification failed: {output}")
    validate_fit(staged)
    print(f"Verified complete staged image: {output}")

    if args.apply:
        backup = apply_to_device(source, original, staged, args.backup_dir)
        print(f"Apply complete. Recovery backup: {backup}")
    else:
        print("Dry-run complete. No block device was written.")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except (InjectionError, OSError) as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        raise SystemExit(1)