#!/bin/sh

# Common helper functions for Raspberry Pi provisioning scripts
# This file should be sourced by all provisioning scripts

# Base directories for various operations.
# Scripts may override these before sourcing this file; the := pattern
# preserves any pre-set value.
: "${LOCK_BASE:=/var/lock/rpi-sb-provisioner}"
: "${LOG_BASE:=/var/log/rpi-sb-provisioner}"
: "${TEMP_BASE:=/srv/rpi-sb-provisioner/tmp}"

# Resource limits
MAX_TEMP_DIR_AGE_HOURS=24

# Creates a temporary directory inside TEMP_BASE so that all temp
# directories are isolated from persistent data (images, workdir, etc.)
# and can be safely reaped by cleanup_orphans.
#
# Usage:
#   make_temp_dir                       # anonymous temp dir
#   make_temp_dir "rpi-sb-bootstrap.XXX" # named template
make_temp_dir() {
    mkdir -p "$TEMP_BASE"
    if [ -n "${1:-}" ]; then
        mktemp -d "$1" --tmpdir="$TEMP_BASE"
    else
        mktemp -d -p "$TEMP_BASE"
    fi
}

# Creates a directory atomically, ensuring no race conditions
# Parameters:
#   $1 - Directory path to create
# Returns:
#   0 - Directory created successfully
#   1 - Directory already exists or creation failed
atomic_mkdir() {
    dir="$1"
    mkdir -p "$(dirname "$dir")"
    if mkdir "$dir" 2>/dev/null; then
        return 0
    fi
    return 1
}

# Executes a command with a file lock
# Parameters:
#   $1 - Lock file path
#   $2 - Timeout in seconds (default: 10)
#   $3 - Command to execute
# Returns:
#   0 - Command executed successfully
#   1 - Failed to acquire lock or command failed
with_lock() {
    lock_file="$1"
    timeout="${2:-10}"
    shift 2

    # /var/lock is a symlink to /run/lock (tmpfs) on Debian, so LOCK_BASE
    # (created in postinst) is wiped on reboot. Recreate the parent dir
    # before touch so with_lock survives a power cycle without reinstall.
    mkdir -p "$(dirname "$lock_file")"

    # Create lock file if it doesn't exist
    touch "$lock_file"
    
    # Use file descriptor 9 for locking
    # Redirecting within a subshell to ensure cleanup
    (
        # Try to acquire lock with timeout
        start_time=$(date +%s)
        end_time=$((start_time + timeout))
        
        while true; do
            if flock -n -x 9; then
                # Lock acquired, run the command
                "$@"
                ret=$?
                # Lock is released when fd 9 is closed upon subshell exit
                exit $ret
            fi
            
            # Check if timeout has been reached
            current_time=$(date +%s)
            if [ "$current_time" -ge "$end_time" ]; then
                # Timeout reached
                exit 1
            fi
            
            # Wait a bit before trying again
            sleep 0.1 2>/dev/null || sleep 1
        done
    ) 9>"$lock_file"
    
    # Return the exit status of the subshell
    return $?
}

# Cleans up orphaned resources (temp dirs older than MAX_TEMP_DIR_AGE_HOURS)
cleanup_orphans() {
    mkdir -p "$TEMP_BASE" 2>/dev/null || true

    # Reap stale temp directories.  TEMP_BASE points to a dedicated tmp
    # subdirectory (/srv/rpi-sb-provisioner/tmp) so persistent data
    # (images, workdir, databases, etc.) is never at risk.
    find "$TEMP_BASE" -mindepth 1 -maxdepth 1 -type d \
        -mmin +"$((MAX_TEMP_DIR_AGE_HOURS * 60))" \
        -exec rm -rf {} + 2>/dev/null || true

    # -mmin +5: skip directories younger than 5 minutes.  Without this gate
    # a concurrent bootstrap@.service that fails fast (e.g. lock contention)
    # racing through cleanup_orphans will reap the empty metadata/ directory
    # of an in-progress sibling bootstrap *between* its mkdir and rpiboot's
    # OTP write -- causing rpiboot's fopen for <serial>.json to ENOENT and
    # the board_type to never reach state.db.
    find "$LOG_BASE" -type d -empty -mmin +5 -delete 2>/dev/null || true
    find "$LOCK_BASE" -type f -mtime +1 -delete 2>/dev/null || true
}

# Resolve the on-disk path for a special flag, accounting for the fact that a
# device's serial representation changes across phases. Pre-triage we typically
# only have the truncated 32-bit ROM serial (8 hex chars, lowercase) or a USB
# path; post-triage we have the full 64-bit DUID (16 hex chars), whose lower
# 32 bits equal the short serial. A flag toggled in one phase must still be
# honoured by the other, so when TARGET_DEVICE_SERIAL is a 16-hex-char DUID we
# also probe the lowercased 8-char suffix.
#
# Arguments:
#   $1 - Special flag directory (e.g. /etc/rpi-sb-provisioner/special-skip-eeprom)
#
# Output (stdout): full path of the matching flag file, or empty if none.
# Exit:   0 if a file was found, 1 otherwise. Safe under `set -e` via `|| true`.
resolve_special_flag_path() {
    _flag_dir="$1"
    _primary="${_flag_dir}/${TARGET_DEVICE_SERIAL}"
    if [ -f "${_primary}" ]; then
        printf '%s\n' "${_primary}"
        return 0
    fi
    case "${TARGET_DEVICE_SERIAL}" in
        [0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F][0-9a-fA-F])
            _short=$(printf '%s' "${TARGET_DEVICE_SERIAL}" | tr 'A-F' 'a-f' | cut -c9-16)
            _fallback="${_flag_dir}/${_short}"
            if [ -f "${_fallback}" ]; then
                printf '%s\n' "${_fallback}"
                return 0
            fi
            ;;
    esac
    return 1
}

announce_start() {
    log "================================================================================"

    log "Starting $1"

    log "================================================================================"
}

announce_stop() {
    log "================================================================================"

    log "Stopping $1"

    log "================================================================================"
}

# Run a command, folding anything it writes to stderr into the caller's log.
#
# Signing is performed by wrapper scripts that rpi-eeprom-digest and
# rpi-sign-bootcode invoke as subprocesses, so their diagnostics (rpi-sb-keyhelper's
# "cannot access PKCS#11 key", "failed to unwrap stored PKCS#11 PIN", ...) go to
# the unit's stderr and land in the journal - not in the per-device log the
# operator is actually reading. Without them a signing failure shows up only as
# an unexplained abort, or worse as a silently unsigned artefact (issue #337).
#
# Returns the command's exit status. errexit is suspended for the command itself
# (it runs as an `if` condition) so stderr is always logged before the caller
# decides what to do with the failure.
log_stderr() {
    _ls_err="$(mktemp)" || { "$@"; return $?; }

    if "$@" 2>"${_ls_err}"; then
        _ls_rc=0
    else
        _ls_rc=$?
    fi

    # `|| [ -n ... ]` so a final line without a trailing newline is not dropped.
    while IFS= read -r _ls_line || [ -n "${_ls_line}" ]; do
        log "${_ls_line}"
    done < "${_ls_err}"
    rm -f "${_ls_err}"

    return "${_ls_rc}"
}

read_config() {
    # Source package defaults first
    if [ -f /usr/share/rpi-sb-provisioner/defaults/config ]; then
        # shellcheck disable=SC1091
        . /usr/share/rpi-sb-provisioner/defaults/config
    else
        printf "%s\n" "Failed to load package defaults. Package may be corrupted." >&2
        return 1
    fi
    
    # Source user overrides (optional - user config takes precedence)
    if [ -f /etc/rpi-sb-provisioner/config ]; then
        # shellcheck disable=SC1091
        . /etc/rpi-sb-provisioner/config
    fi
}

# Initializes fastboot connection and device identification variables
# Takes a fastboot device specifier (USB serial or network address) and:
# 1. Verifies fastboot connectivity
# 2. Gets the device serial number
# 3. Discovers whether the daemon advertises a data-plane-only TCP path
# 4. Tests both IPv4 and IPv6 connectivity, preferring IPv6 if available
# 5. Determines the USB path for the device
#
# Arguments:
#   $1 - Fastboot device specifier (required, USB serial)
#
# Sets the following global variables:
#   FASTBOOT_DEVICE_SPECIFIER       - Control-plane specifier (USB serial,
#                                     unless the daemon is in legacy TCP-full
#                                     mode and a TCP route is reachable, in
#                                     which case it is promoted to tcp:<addr>)
#   FASTBOOT_TCP_FLASH_SPECIFIER    - Set to tcp:<addr> when the daemon
#                                     advertises tcp-data-plane-only=yes and a
#                                     TCP route is reachable; empty otherwise.
#                                     Callers should use this for `flash` and
#                                     fall back to FASTBOOT_DEVICE_SPECIFIER.
#   TARGET_DEVICE_SERIAL            - Device serial number
#   TARGET_USB_PATH                 - USB device path (e.g., 1-1.2)
#   TARGET_DEVICE_PATH              - USB device node (e.g., /dev/bus/usb/001/004)
#
# Exits with error if:
#   - Cannot establish fastboot connection
#   - Cannot determine USB path
setup_fastboot_and_id_vars() {
    FASTBOOT_DEVICE_SPECIFIER="$1"
    FASTBOOT_TCP_FLASH_SPECIFIER=""
    export FASTBOOT_TCP_FLASH_SPECIFIER

    timeout_fatal fastboot -s "${FASTBOOT_DEVICE_SPECIFIER}" getvar version
    TARGET_DEVICE_SERIAL="$(get_variable serialno)"

    # Returns "yes" if the device-side fastbootd is running in -i usb+tcp
    # split mode. Empty/missing/"no" means legacy behaviour (TCP, when
    # reachable, carries the full command stream).
    TCP_DATA_PLANE_ONLY="$(get_variable tcp-data-plane-only)"

    announce_start "Testing Fastboot IP connectivity"
    USE_IPV4=
    USE_IPV6=
    set +e
    IPV6_ADDRESS="$(fastboot -s "${FASTBOOT_DEVICE_SPECIFIER}" getvar ipv6-address 2>&1 | awk '/^ipv6-address:/ {print $2}')"
    (timeout_nonfatal fastboot -s tcp:"${IPV6_ADDRESS}" getvar version)
    USE_IPV6=$?
    IPV4_ADDRESS="$(fastboot -s "${FASTBOOT_DEVICE_SPECIFIER}" getvar ipv4-address 2>&1 | grep -oE '([0-9]{1,3}\.){3}[0-9]{1,3}')"
    (timeout_nonfatal fastboot -s tcp:"${IPV4_ADDRESS}" getvar version)
    USE_IPV4=$?
    set -e

    # Pick the best reachable TCP specifier (IPv6 preferred).
    if [ "${USE_IPV6}" -eq 0 ]; then
        TCP_SPEC="tcp:${IPV6_ADDRESS}"
    elif [ "${USE_IPV4}" -eq 0 ]; then
        TCP_SPEC="tcp:${IPV4_ADDRESS}"
    else
        TCP_SPEC=""
    fi

    if [ "${TCP_DATA_PLANE_ONLY}" = "yes" ]; then
        # Split mode: keep control plane on USB and route flash over TCP if
        # we have a route. Falling back to USB-only flashing is fine — the
        # daemon's USB FastbootDevice carries the full command set.
        FASTBOOT_DEVICE_SPECIFIER="${TARGET_DEVICE_SERIAL}"
        FASTBOOT_TCP_FLASH_SPECIFIER="${TCP_SPEC}"
        if [ -n "${FASTBOOT_TCP_FLASH_SPECIFIER}" ]; then
            log "Split-mode daemon: control over USB ${FASTBOOT_DEVICE_SPECIFIER}, flash over ${FASTBOOT_TCP_FLASH_SPECIFIER}"
        else
            log "Split-mode daemon advertised but no TCP route reachable; flash will fall back to USB"
        fi
    else
        # Legacy daemon: promote everything to TCP when reachable.
        if [ -n "${TCP_SPEC}" ]; then
            FASTBOOT_DEVICE_SPECIFIER="${TCP_SPEC}"
        else
            FASTBOOT_DEVICE_SPECIFIER="${TARGET_DEVICE_SERIAL}"
        fi
    fi

    # Set TARGET_USB_PATH and TARGET_DEVICE_PATH based on TARGET_DEVICE_SERIAL
    if [ -n "${TARGET_DEVICE_SERIAL}" ]; then
        # Try to get the USB path for the device
        usb_path=$(get_usb_path_for_serial "${TARGET_DEVICE_SERIAL}")
        
        if [ -n "$usb_path" ]; then
            TARGET_USB_PATH="$usb_path"
            log "Found USB path ${TARGET_USB_PATH} for device ${TARGET_DEVICE_SERIAL}"
        else
            log "Warning: Could not find USB path for device ${TARGET_DEVICE_SERIAL}"

        fi

        device_path=$(get_device_path_for_serial "${TARGET_DEVICE_SERIAL}" || true)
        if [ -n "$device_path" ]; then
            TARGET_DEVICE_PATH="$device_path"
            log "Found device path ${TARGET_DEVICE_PATH} for device ${TARGET_DEVICE_SERIAL}"
        else
            log "Warning: Could not find device path for device ${TARGET_DEVICE_SERIAL}"
        fi
    fi

    # Ensure TARGET_USB_PATH is set
    if [ -z "${TARGET_USB_PATH}" ]; then
        log "Error: Could not determine USB path for device ${TARGET_DEVICE_SERIAL}"
        record_state "${TARGET_DEVICE_SERIAL}" "${PROVISIONER_ABORTED}" "unknown-usb-path"
        exit 1
    fi
}

# Export device-identity environment variables for customisation hooks.
# Positional arguments remain the contract for existing per-stage fields;
# newly exposed identity fields (USB path, device path on provisioning hooks)
# are provided via the environment so scripts can share logic across stages.
export_customisation_env() {
    export TARGET_USB_PATH="${TARGET_USB_PATH:-}"
    export TARGET_DEVICE_PATH="${TARGET_DEVICE_PATH:-}"
    export TARGET_DEVICE_SERIAL="${TARGET_DEVICE_SERIAL:-}"
    export TARGET_DEVICE_FAMILY="${TARGET_DEVICE_FAMILY:-}"
    export FASTBOOT_DEVICE_SPECIFIER="${FASTBOOT_DEVICE_SPECIFIER:-}"
    export RPI_DEVICE_STORAGE_TYPE="${RPI_DEVICE_STORAGE_TYPE:-}"
}

# Echo the customisation-script provisioner prefix for this station's
# configured image route (IDP directory vs traditional PROVISIONING_STYLE).
get_configured_provisioner_name() {
    if [ -d "${GOLD_MASTER_OS_FILE}" ]; then
        echo "idp-provisioner"
        return 0
    fi
    case ${PROVISIONING_STYLE} in
        "secure-boot")
            echo "sb-provisioner"
            ;;
        "fde-only")
            echo "fde-provisioner"
            ;;
        "naked")
            echo "naked-provisioner"
            ;;
        *)
            echo ""
            ;;
    esac
}

# Map a storage type to its Raspberry Pi BOOT_ORDER nibble.
# See https://www.raspberrypi.com/documentation/computers/raspberry-pi.html#BOOT_ORDER
# Accepts both the raw config values (sd/emmc/nvme) and the resolved block
# device names (mmcblk0/nvme0n1) for robustness. Returns non-zero for anything
# unmapped so callers can skip cleanly rather than corrupt the boot order.
storage_type_to_boot_code() {
    case "$1" in
        sd|emmc|mmcblk0) echo "1" ;;   # SD/eMMC share the SDIO boot path
        nvme|nvme0n1)    echo "6" ;;   # NVMe
        *)               return 1 ;;
    esac
}

# Resolve the storage type (sd/emmc/nvme) to use for boot-order decisions.
# Prefers the station config, then falls back to the IDP artefact's declared
# storage type. The fallback is what makes boot-order matching work for IDP
# runs during bootstrap, where RPI_DEVICE_STORAGE_TYPE is often unset because
# the IDP provisioner only adopts it from the JSON later in the pipeline.
# Echoes the resolved type, or nothing if it cannot be determined.
resolve_effective_storage_type() {
    if [ -n "${RPI_DEVICE_STORAGE_TYPE:-}" ]; then
        printf '%s\n' "${RPI_DEVICE_STORAGE_TYPE}"
        return 0
    fi

    if [ -n "${GOLD_MASTER_OS_FILE:-}" ] && [ -d "${GOLD_MASTER_OS_FILE}" ] \
       && command -v jq >/dev/null 2>&1; then
        _idp_json=$(find "${GOLD_MASTER_OS_FILE}" -maxdepth 1 -name '*.json' -type f 2>/dev/null | head -n1)
        if [ -n "${_idp_json}" ]; then
            _st=$(jq -r '.IGmeta.IGconf_device_storage_type // empty' < "${_idp_json}" 2>/dev/null)
            [ -n "${_st}" ] && printf '%s\n' "${_st}"
        fi
    fi

    # Always succeed: an undetermined storage type is a caller-handled
    # condition (empty output), not a failure that should trip `set -e`.
    return 0
}

# Move hex nibble $2 to the least-significant (first-tried) position of the
# BOOT_ORDER digit string $1, preserving the remaining fallbacks in order.
# Idempotent: promoting an already-first nibble yields the same string. If the
# nibble is absent it is appended, adding the device as the first boot option.
# e.g. _promote_boot_nibble "f2461" 6 -> "f2416"
_promote_boot_nibble() {
    _digits="$1"
    _target="$2"
    _out=""
    _removed=0
    while [ -n "${_digits}" ]; do
        _c=${_digits%"${_digits#?}"}   # first char
        _digits=${_digits#?}           # rest
        if [ "${_removed}" -eq 0 ] && [ "${_c}" = "${_target}" ]; then
            _removed=1
            continue
        fi
        _out="${_out}${_c}"
    done
    # BOOT_ORDER is tried least-significant-nibble first, i.e. right-to-left,
    # so the promoted device must be the rightmost (last) digit.
    printf '%s%s\n' "${_out}" "${_target}"
}

# Whether storage-based BOOT_ORDER matching is enabled. This feature is on by
# default; it is only disabled when explicitly set to an empty or falsey value.
# The WebUI toggle stores "1" (on) / "" (off); hand-edited configs may also use
# 0/false/no/off. An entirely unset variable falls back to on (default-on).
boot_order_match_storage_enabled() {
    # Single-dash default: an unset variable falls back to on (default-on),
    # but an explicitly empty value (how the WebUI stores "off") disables it.
    case "${RPI_DEVICE_BOOT_ORDER_MATCH_STORAGE-1}" in
        ''|0|false|FALSE|False|no|NO|No|off|OFF|Off) return 1 ;;
        *) return 0 ;;
    esac
}

# When storage-based BOOT_ORDER matching is enabled, rewrite the BOOT_ORDER
# line in the given bootloader config file so the target storage device is the
# first boot option, keeping the shipped fallbacks (and restart) behind it.
#
# Must be called before the config is signed/flashed. No-op (and never fatal)
# when: the feature is disabled, the storage type can't be resolved or isn't
# mappable, the config file is missing, or it doesn't contain exactly one
# BOOT_ORDER line (we leave hand-authored multi-section configs untouched).
maybe_reorder_boot_order_for_storage() {
    _config="$1"

    boot_order_match_storage_enabled || return 0

    if [ ! -f "${_config}" ]; then
        log "BOOT_ORDER reorder: config file '${_config}' not found; skipping"
        return 0
    fi

    _stype=$(resolve_effective_storage_type)
    if [ -z "${_stype}" ]; then
        log "BOOT_ORDER reorder requested but storage type is unknown; leaving BOOT_ORDER unchanged"
        return 0
    fi

    if ! _code=$(storage_type_to_boot_code "${_stype}"); then
        log "BOOT_ORDER reorder: unsupported storage type '${_stype}'; leaving BOOT_ORDER unchanged"
        return 0
    fi

    _count=$(grep -cE '^BOOT_ORDER=' "${_config}" 2>/dev/null || true)
    if [ "${_count}" != "1" ]; then
        log "BOOT_ORDER reorder: expected exactly one BOOT_ORDER line in '${_config}', found ${_count}; leaving config untouched"
        return 0
    fi

    _val=$(grep -E '^BOOT_ORDER=' "${_config}" | sed 's/^BOOT_ORDER=//')
    _hex=${_val#0x}
    _hex=${_hex#0X}
    _new=$(_promote_boot_nibble "${_hex}" "${_code}")

    if [ "0x${_new}" = "${_val}" ]; then
        log "BOOT_ORDER already prioritises storage '${_stype}' (${_val}); no change"
        return 0
    fi

    sed -i -e "s/^BOOT_ORDER=.*/BOOT_ORDER=0x${_new}/" "${_config}"
    log "BOOT_ORDER reordered for storage '${_stype}': ${_val} -> 0x${_new}"
}

# Run the provision-failed hook for programming-rig signalling (e.g. LEDs).
# HOOK_CONTEXT is "bootstrap" (serial/family/usb/device path args) or
# "provisioning" (fastboot specifier/serial/storage args).
# Mark the failure in progress as a permanent misconfiguration -- something
# retrying the same device will never resolve, such as no OS image or no
# signing key being configured. Distinct from the stage the failure happened
# in, which run_provision_failed_hook already reports: a bootstrap-stage
# failure may be either transient (USB re-enumeration churn as the device
# reboots) or permanent, and hooks that want to react loudly to the latter
# cannot tell them apart from the stage alone.
mark_permanent_failure() {
    PROVISION_FAILURE_PERMANENT=1
    export PROVISION_FAILURE_PERMANENT
}

# Classify the configured OS image, failing early when it cannot be used.
#
# GOLD_MASTER_OS_FILE names either a regular file (a whole-disk .img) or a
# directory (an rpi-image-gen IDP artefact). Telling "not configured",
# "configured but absent" and "present but unusable" apart matters because
# triage picks the provisioner from the *shape* of this path: an IDP artefact
# whose directory is momentarily missing -- an unmounted share, say -- is
# otherwise indistinguishable from a .img, and gets routed to a provisioner
# that cannot read it and fails much later with an unrelated error.
#
# Results come back in globals, deliberately, and nothing is written to
# stdout: capturing a reason with $( ) would run this in a subshell and throw
# GOLD_MASTER_OS_KIND away with it. Call it directly:
#
#     if ! classify_gold_master_os; then
#         die "${GOLD_MASTER_OS_ERROR}"
#     fi
#
# Sets and exports GOLD_MASTER_OS_KIND to "image" or "idp" on success, so
# callers -- and the customisation hooks they run -- can branch on the form of
# the artefact without re-testing the path. Sets GOLD_MASTER_OS_ERROR to the
# reason on failure.
# Exit:   0 when usable, 1 otherwise.
classify_gold_master_os() {
    GOLD_MASTER_OS_KIND=""
    GOLD_MASTER_OS_ERROR=""
    export GOLD_MASTER_OS_KIND GOLD_MASTER_OS_ERROR

    if [ -z "${GOLD_MASTER_OS_FILE}" ]; then
        GOLD_MASTER_OS_ERROR="No OS image is configured. Set GOLD_MASTER_OS_FILE in /etc/rpi-sb-provisioner/config, or select an image in the web UI."
        return 1
    fi

    if [ -d "${GOLD_MASTER_OS_FILE}" ]; then
        # IDP artefact. Apply the same descriptor rule the IDP provisioner
        # enforces in its own pre-flight, so the two cannot disagree about
        # whether a directory is a usable artefact.
        _cgm_json_count=$(find "${GOLD_MASTER_OS_FILE}" -maxdepth 1 -name '*.json' -type f 2>/dev/null | wc -l)
        if [ "${_cgm_json_count}" -eq 0 ]; then
            GOLD_MASTER_OS_ERROR="GOLD_MASTER_OS_FILE is a directory with no JSON description file: ${GOLD_MASTER_OS_FILE}. An IDP artefact directory must contain exactly one."
            return 1
        fi
        if [ "${_cgm_json_count}" -gt 1 ]; then
            GOLD_MASTER_OS_ERROR="GOLD_MASTER_OS_FILE is a directory with ${_cgm_json_count} JSON description files: ${GOLD_MASTER_OS_FILE}. An IDP artefact directory must contain exactly one."
            return 1
        fi
        GOLD_MASTER_OS_KIND="idp"
        return 0
    fi

    if [ -f "${GOLD_MASTER_OS_FILE}" ]; then
        if [ ! -s "${GOLD_MASTER_OS_FILE}" ]; then
            GOLD_MASTER_OS_ERROR="Configured OS image is empty: ${GOLD_MASTER_OS_FILE}"
            return 1
        fi
        if [ ! -r "${GOLD_MASTER_OS_FILE}" ]; then
            GOLD_MASTER_OS_ERROR="Configured OS image is not readable: ${GOLD_MASTER_OS_FILE}"
            return 1
        fi
        GOLD_MASTER_OS_KIND="image"
        return 0
    fi

    if [ -e "${GOLD_MASTER_OS_FILE}" ]; then
        GOLD_MASTER_OS_ERROR="Configured OS image is neither a regular file nor a directory: ${GOLD_MASTER_OS_FILE}"
        return 1
    fi

    # Absent entirely. Deliberately does not guess whether a .img or an IDP
    # directory was intended -- the path is gone, so the shape is unknowable,
    # and a mount that has not come back is the common cause.
    GOLD_MASTER_OS_ERROR="Configured OS image does not exist: ${GOLD_MASTER_OS_FILE}. If this is an IDP artefact directory on a removable or network mount, check the mount is present."
    return 1
}

run_provision_failed_hook() {
    PROVISIONER_NAME="$1"
    HOOK_CONTEXT="${2:-provisioning}"

    if [ -z "${PROVISIONER_NAME}" ]; then
        log "No provisioner configured; skipping provision-failed hook"
        return 0
    fi

    set +e
    # Reported alongside the stage rather than folded into it: the stage still
    # selects the hook's positional arguments, so widening it would change the
    # contract every existing hook is written against.
    export PROVISION_FAILED_PERMANENT="${PROVISION_FAILURE_PERMANENT:-0}"
    if [ "${HOOK_CONTEXT}" = "bootstrap" ]; then
        export PROVISION_FAILED_CONTEXT=bootstrap
        run_customisation_script "${PROVISIONER_NAME}" "provision-failed" \
            "${TARGET_DEVICE_SERIAL}" "${TARGET_DEVICE_FAMILY}" \
            "${TARGET_USB_PATH}" "${TARGET_DEVICE_PATH}"
    else
        export PROVISION_FAILED_CONTEXT=provisioning
        run_customisation_script "${PROVISIONER_NAME}" "provision-failed" \
            "${FASTBOOT_DEVICE_SPECIFIER}" "${TARGET_DEVICE_SERIAL}" \
            "${RPI_DEVICE_STORAGE_TYPE}"
    fi
    unset PROVISION_FAILED_CONTEXT
    unset PROVISION_FAILED_PERMANENT
}

run_customisation_script() {
    PROVISIONER_NAME="$1"
    STAGE_NAME="$2"
    
    SCRIPT_NAME="${PROVISIONER_NAME}-${STAGE_NAME}.sh"
    SCRIPT_PATH="/etc/rpi-sb-provisioner/scripts/${SCRIPT_NAME}"
    
    if [ -f "${SCRIPT_PATH}" ]; then
        announce_start "Running customisation script for ${PROVISIONER_NAME} at stage ${STAGE_NAME}"
        
        # Check if script is executable
        if [ ! -x "${SCRIPT_PATH}" ]; then
            log "Skipping disabled customisation script: ${SCRIPT_NAME}"
            return 0
        fi
        
        # Temporarily disable error exit so we can capture the script's exit code
        # Save current error exit setting and disable it
        ERROR_EXIT_WAS_SET=0
        case $- in
            *e*) ERROR_EXIT_WAS_SET=1 ;;
        esac
        set +e
        export_customisation_env
        if [ "${STAGE_NAME}" = "post-flash" ]; then
            export_manufacturing_env
        fi
        # Handle different stages with appropriate parameters
        if [ "${STAGE_NAME}" = "post-flash" ]; then
            # For post-flash stage, pass device info that can be used with fastboot.
            # TARGET_USB_PATH and TARGET_DEVICE_PATH are also exported (issue #273).
            "${SCRIPT_PATH}" "${FASTBOOT_DEVICE_SPECIFIER}" "${TARGET_DEVICE_SERIAL}" "${RPI_DEVICE_STORAGE_TYPE}"
        elif [ "${STAGE_NAME}" = "provision-started" ]; then
            # For provision-started stage, forward (fastboot specifier, serial,
            # storage type) so hooks can signal programming rigs (e.g. LED).
            # TARGET_USB_PATH and TARGET_DEVICE_PATH are also exported (issue #273).
            FASTBOOT_SPEC_ARG="$3"
            TARGET_SERIAL_ARG="$4"
            STORAGE_TYPE_ARG="$5"
            "${SCRIPT_PATH}" "${FASTBOOT_SPEC_ARG}" "${TARGET_SERIAL_ARG}" "${STORAGE_TYPE_ARG}"
        elif [ "${STAGE_NAME}" = "provision-failed" ]; then
            if [ "${PROVISION_FAILED_CONTEXT:-provisioning}" = "bootstrap" ]; then
                "${SCRIPT_PATH}" "$3" "$4" "$5" "$6"
            else
                "${SCRIPT_PATH}" "$3" "$4" "$5"
            fi
        elif [ "${STAGE_NAME}" = "bootstrap" ]; then
            # For bootstrap stage, pass device detection info
            TARGET_DEVICE_SERIAL_ARG="$3"
            TARGET_DEVICE_FAMILY_ARG="$4"
            TARGET_USB_PATH_ARG="$5"
            TARGET_DEVICE_PATH_ARG="$6"
            "${SCRIPT_PATH}" "${TARGET_DEVICE_SERIAL_ARG}" "${TARGET_DEVICE_FAMILY_ARG}" "${TARGET_USB_PATH_ARG}" "${TARGET_DEVICE_PATH_ARG}"
        else
            # For filesystem mount stages, pass mount points
            BOOT_MOUNT="$3"
            ROOTFS_MOUNT="$4"
            "${SCRIPT_PATH}" "${BOOT_MOUNT}" "${ROOTFS_MOUNT}"
        fi
        # Capture exit code immediately, before restoring set -e
        SCRIPT_EXIT_CODE=$?
        if [ "${ERROR_EXIT_WAS_SET}" -eq 1 ]; then
            set -e
        fi
        
        if [ $SCRIPT_EXIT_CODE -eq 0 ]; then
            announce_stop "Customisation script ${SCRIPT_NAME} completed successfully"
        else
            announce_stop "Customisation script ${SCRIPT_NAME} failed with exit code ${SCRIPT_EXIT_CODE}"
            log "ERROR: Customisation script ${SCRIPT_NAME} failed with exit code ${SCRIPT_EXIT_CODE}"
            return $SCRIPT_EXIT_CODE
        fi
    else
        log "No customisation script found for ${PROVISIONER_NAME} at stage ${STAGE_NAME}"
    fi
}

# Retrieve a single-line fastboot variable. NOTE: the value is truncated at
# the first newline -- use get_variable_pem() for multi-line values such as
# public-key/private-key, which the gadget returns as armoured PEM.
get_variable() {
    fastboot -s "${FASTBOOT_DEVICE_SPECIFIER}" getvar "$1" 2>&1 | grep -oP "${1}"': \K[^\r\n]*' || true
}

# Retrieve a fastboot variable whose value is an armoured PEM block, printing
# the whole block (BEGIN line through END line) on stdout.
#
# rpi-fastbootd returns public-key/private-key as multi-line PEM, so
# get_variable() captures only "-----BEGIN PUBLIC KEY-----" and drops the
# body. Extract the armoured block instead: strip fastboot's "<var>: " prefix
# from the BEGIN line and ignore its trailing progress output.
#
# Returns 1 (printing nothing) when the device did not return a PEM block --
# e.g. private-key answers "refused" once the OTP ECDSA key-slot is LOCKED,
# which rpi-fastbootd does at startup and again immediately after
# `oem fwcrypto init`.
get_variable_pem() {
    _gvp_output=$(fastboot -s "${FASTBOOT_DEVICE_SPECIFIER}" getvar "$1" 2>&1 | tr -d '\r') || true
    _gvp_pem=$(printf '%s\n' "${_gvp_output}" | \
        sed -n '/-----BEGIN /,/-----END /{
            s/^.*\(-----BEGIN \)/\1/
            p
        }')
    if [ -z "${_gvp_pem}" ]; then
        return 1
    fi
    printf '%s\n' "${_gvp_pem}"
}

# =============================================================================
# Loop Device and Mount Helpers
# =============================================================================
# These functions manage loop devices and mounts for OS image manipulation.
# Lifted from pi-gen/scripts/common.
# =============================================================================

ensure_next_loopdev() {
    set +e
    loopdev="$(losetup -f)"
    loopmaj="$(echo "$loopdev" | sed -E 's/.*[0-9]*?([0-9]+)$/\1/')"
    [ -b "$loopdev" ] || mknod "$loopdev" b 7 "$loopmaj"
    set -e
}

ensure_loopdev_partitions() {
    set +e
    lsblk -r -n -o "NAME,MAJ:MIN" "$1" | grep -v "^${1#/dev/}" | while read -r line; do
        partition="${line%% *}"
        majmin="${line#* }"
        if [ ! -b "/dev/$partition" ]; then
            mknod "/dev/$partition" b "${majmin%:*}" "${majmin#*:}"
        fi
    done
    set -e
}

unmount() {
    if [ -z "$1" ]; then
        DIR=$PWD
    else
        DIR=$1
    fi

    while mount | grep -q "$DIR"; do
        locs=$(mount | grep "$DIR" | cut -f 3 -d ' ' | sort -r)
        for loc in $locs; do
            umount "$loc"
        done
    done
}

unmount_image() {
    sync
    sleep 1
    losetup --list | grep "$1" | cut -f1 -d' ' | while read -r LOOP_DEVICE; do
        [ -z "$LOOP_DEVICE" ] && continue
        for part in "${LOOP_DEVICE}"p*; do
            if DIR=$(findmnt -n -o target -S "$part"); then
                unmount "$DIR"
            fi
        done
        losetup -d "$LOOP_DEVICE"
    done
}

# =============================================================================
# Signing Infrastructure
# =============================================================================
# These functions provide a unified interface for cryptographic signing
# operations, supporting both PEM files and PKCS#11 hardware security modules.
#
# Usage:
#   1. Call init_signing_context() early in your script (after read_config)
#   2. Use the get_*_sign_args() functions to get tool-specific arguments
#
# Supported signing modes:
#   - "pem"    : Traditional PEM file-based signing
#   - "pkcs11" : PKCS#11 HSM-based signing via wrapper script
#   - "none"   : No signing key configured
# =============================================================================

SIGNING_MODE=""
CUSTOMER_PUBLIC_KEY_FILE=""
PKCS11_WRAPPER_SCRIPT="/usr/bin/rpi-sb-pkcs11-sign.sh"
PEM_WRAPPER_SCRIPT="/usr/bin/rpi-sb-pem-sign.sh"

# Initialize the signing context
# Must be called after read_config() and before any signing operations
#
# This function:
#   - Validates that only one key source is configured (mutual exclusion)
#   - Derives and caches the public key (needed for rpi-sign-bootcode -p)
#   - Sets SIGNING_MODE for use by other functions
#
# Exits with error if:
#   - Both CUSTOMER_KEY_FILE_PEM and CUSTOMER_KEY_PKCS11_NAME are set
#   - PKCS#11 key is not accessible
#   - PEM key file doesn't exist
init_signing_context() {
    OPENSSL="${OPENSSL:-openssl}"
    
    # Mutual exclusion check
    if [ -n "${CUSTOMER_KEY_PKCS11_NAME}" ] && [ -n "${CUSTOMER_KEY_FILE_PEM}" ]; then
        log "ERROR: Both CUSTOMER_KEY_PKCS11_NAME and CUSTOMER_KEY_FILE_PEM are set."
        log "Please configure only one signing key source."
        return 1
    fi
    
    if [ -n "${CUSTOMER_KEY_PKCS11_NAME}" ]; then
        SIGNING_MODE="pkcs11"
        log "Signing mode: PKCS#11 (${CUSTOMER_KEY_PKCS11_NAME})"
        
        # Verify the PKCS#11 key is accessible and derive its public key. This
        # goes through rpi-sb-keyhelper, which supplies the stored (device-
        # wrapped) token PIN to the pkcs11-provider in-process. Some tokens
        # require a C_Login even to read the public key, so the previous bare
        # `openssl pkey` call failed here whenever the PIN lived only in the
        # WebUI's pkcs11.pin store (issue #328).
        CUSTOMER_PUBLIC_KEY_FILE="$(mktemp)"
        if ! rpi-sb-keyhelper pubkey --pkcs11-uri "${CUSTOMER_KEY_PKCS11_NAME}" \
                > "${CUSTOMER_PUBLIC_KEY_FILE}" 2>/dev/null; then
            rm -f "${CUSTOMER_PUBLIC_KEY_FILE}"
            CUSTOMER_PUBLIC_KEY_FILE=""
            log "ERROR: Cannot access PKCS#11 key: ${CUSTOMER_KEY_PKCS11_NAME}"
            log "Check that the token is present, PIN is correct, and URI is valid."
            return 1
        fi
        log "PKCS#11 key validated and public key derived"
        
    elif [ -n "${CUSTOMER_KEY_FILE_PEM}" ]; then
        log "Signing mode: PEM file (${CUSTOMER_KEY_FILE_PEM})"

        if [ ! -f "${CUSTOMER_KEY_FILE_PEM}" ]; then
            log "ERROR: PEM key file not found: ${CUSTOMER_KEY_FILE_PEM}"
            return 1
        fi

        # A PEM key may be stored device-wrapped at rest (magic "RPISBW1" in
        # the first bytes). A wrapped key is never unwrapped here: signing goes
        # through rpi-sb-keyhelper so the plaintext key stays inside that
        # process. A legacy plaintext key keeps the original openssl path.
        if [ "$(head -c 7 "${CUSTOMER_KEY_FILE_PEM}" 2>/dev/null)" = "RPISBW1" ]; then
            SIGNING_MODE="pem-wrapped"
            # Derive the (non-secret) public key via the helper, which unwraps
            # in-memory. Reuses the existing CUSTOMER_PUBLIC_KEY_FILE lifecycle.
            CUSTOMER_PUBLIC_KEY_FILE="$(mktemp)"
            if ! rpi-sb-keyhelper pubkey --key "${CUSTOMER_KEY_FILE_PEM}" \
                    > "${CUSTOMER_PUBLIC_KEY_FILE}" 2>/dev/null; then
                rm -f "${CUSTOMER_PUBLIC_KEY_FILE}"
                CUSTOMER_PUBLIC_KEY_FILE=""
                log "ERROR: Failed to derive public key from wrapped PEM key"
                log "Check this is the board the key was wrapped on, and that the key is not corrupt."
                return 1
            fi
            log "Wrapped PEM key validated and public key derived"
        else
            SIGNING_MODE="pem"
            # Derive public key
            CUSTOMER_PUBLIC_KEY_FILE="$(mktemp)"
            if ! "${OPENSSL}" rsa -in "${CUSTOMER_KEY_FILE_PEM}" -pubout > "${CUSTOMER_PUBLIC_KEY_FILE}" 2>/dev/null; then
                rm -f "${CUSTOMER_PUBLIC_KEY_FILE}"
                CUSTOMER_PUBLIC_KEY_FILE=""
                log "ERROR: Failed to derive public key from: ${CUSTOMER_KEY_FILE_PEM}"
                return 1
            fi
            log "PEM key validated and public key derived"
        fi

    else
        SIGNING_MODE="none"
        log "Signing mode: none (no signing key configured)"
    fi
    
    export SIGNING_MODE CUSTOMER_PUBLIC_KEY_FILE
    return 0
}

# Check if signing is available
# Returns 0 if a signing key is configured, 1 otherwise
signing_available() {
    [ "${SIGNING_MODE}" = "pem" ] || [ "${SIGNING_MODE}" = "pem-wrapped" ] || [ "${SIGNING_MODE}" = "pkcs11" ]
}

# Get arguments for OpenSSL dgst -sign operations
# Usage: openssl dgst -sign $(get_openssl_sign_args) -sha256 <file>
#
# Returns the key specification and format arguments for OpenSSL
# Exits with error if no signing key is configured
get_openssl_sign_args() {
    case "${SIGNING_MODE}" in
        pkcs11)
            # pkcs11-provider resolves the pkcs11: URI via OSSL_STORE; the
            # deprecated -engine/-keyform engine flags are no longer needed.
            echo "${CUSTOMER_KEY_PKCS11_NAME} -provider pkcs11 -provider default"
            ;;
        pem)
            echo "${CUSTOMER_KEY_FILE_PEM} -keyform PEM"
            ;;
        *)
            # NB: pem-wrapped has no openssl directive form (the key is never on
            # disk). Such callers must use sign_image_hex(), not this function.
            mark_permanent_failure
            log "ERROR: No signing key configured (get_openssl_sign_args)"
            return 1
            ;;
    esac
}

# Sign a file and emit the signature as lowercase hex on stdout (newline
# terminated), honouring the active SIGNING_MODE. This is the single place that
# knows how to sign for every mode: a device-wrapped PEM key is signed inside
# rpi-sb-keyhelper (the plaintext key never touches disk or this shell), while
# every other mode uses `openssl dgst` with the directives from
# get_openssl_sign_args. Use this for the common "append rsa2048 hex to a .sig
# sidecar" pattern instead of calling openssl dgst directly.
sign_image_hex() {
    _sign_src="$1"
    case "${SIGNING_MODE}" in
        pem-wrapped)
            "${PEM_WRAPPER_SCRIPT}" -a rsa2048-sha256 "${_sign_src}" || return 1
            echo
            ;;
        pkcs11)
            # Route through the wrapper (not a bare `openssl dgst`) so the
            # stored token PIN is supplied in-process: a direct openssl call
            # here has no way to log in to the token. Wrapper emits hex with no
            # trailing newline, so add one to match the other modes' output.
            "${PKCS11_WRAPPER_SCRIPT}" -a rsa2048-sha256 "${_sign_src}" || return 1
            echo
            ;;
        *)
            # shellcheck disable=SC2046
            "${OPENSSL}" dgst -sign $(get_openssl_sign_args) -sha256 "${_sign_src}" | xxd -c 4096 -p
            ;;
    esac
}

# Assert that a .sig sidecar actually carries a usable RSA-2048 signature.
#
# Every producer of these files builds the signature inside a command
# substitution - `echo "rsa2048: $(signer ...)"` - both in rpi-eeprom-digest
# upstream and (historically) here. Command substitution discards the signer's
# exit status, so when signing fails the hash/ts lines are still written, an
# empty `rsa2048:` line is appended, and the producer exits 0. errexit never
# fires and the result is a structurally plausible file that the device only
# rejects at boot, from RECOVERY, with "Failed to parse signature file" ->
# "Bad signature pieeprom.sig" (issue #337). Check the artefact instead of
# trusting the exit status.
#
# Arguments:
#   $1 - path to the .sig file
#   $2 - optional description used in error messages (defaults to the path)
#
# Returns 0 only if the file contains an rsa2048 line of exactly 512
# hexadecimal characters. The Raspberry Pi boot ROM accepts nothing but
# RSA-2048 (256 bytes) here, so a signature of any other length is unusable
# regardless of how it was produced.
validate_sig_file() {
    _sig_file="$1"
    _sig_desc="${2:-$1}"

    if [ ! -s "${_sig_file}" ]; then
        log "ERROR: signature file is missing or empty: ${_sig_desc}"
        return 1
    fi

    _sig_hex="$(sed -n 's/^rsa2048:[[:space:]]*//p' "${_sig_file}" | tr -d '\r\n')"

    if [ -z "${_sig_hex}" ]; then
        log "ERROR: no rsa2048 signature in ${_sig_desc} - the signing key produced no output."
        log "ERROR: for a PKCS#11 token, check for rpi-sb-keyhelper errors above (token present? PIN saved and correct? URI valid?)."
        return 1
    fi

    if [ "${#_sig_hex}" -ne 512 ]; then
        log "ERROR: malformed rsa2048 signature in ${_sig_desc}: expected 512 hex characters (RSA-2048), got ${#_sig_hex}"
        return 1
    fi

    case "${_sig_hex}" in
        *[!0-9a-fA-F]*)
            log "ERROR: rsa2048 signature in ${_sig_desc} is not hexadecimal"
            return 1
            ;;
    esac

    return 0
}

# Get arguments for rpi-eeprom-digest signing operations
# Usage: rpi-eeprom-digest $(get_eeprom_digest_sign_args) -i <input> -o <output>
#
# Returns -k or -H arguments depending on signing mode
# Returns empty string if no signing (unsigned mode)
get_eeprom_digest_sign_args() {
    case "${SIGNING_MODE}" in
        pkcs11)
            echo "-H ${PKCS11_WRAPPER_SCRIPT}"
            ;;
        pem-wrapped)
            # Wrapped key: sign via the helper so the key is never on disk.
            echo "-H ${PEM_WRAPPER_SCRIPT}"
            ;;
        pem)
            echo "-k ${CUSTOMER_KEY_FILE_PEM}"
            ;;
        *)
            # No signing - return empty for unsigned operation
            echo ""
            ;;
    esac
}

# Get arguments for rpi-sign-bootcode signing operations
# Usage: rpi-sign-bootcode $(get_sign_bootcode_key_args) -c <chip> -i <input> -o <output> ...
#
# Returns -k or -H/-p arguments depending on signing mode
# Exits with error if no signing key is configured (bootcode always requires signing)
get_sign_bootcode_key_args() {
    case "${SIGNING_MODE}" in
        pkcs11)
            # PKCS#11 mode requires HSM wrapper AND public key
            if [ -z "${CUSTOMER_PUBLIC_KEY_FILE}" ]; then
                log "ERROR: Public key not available for rpi-sign-bootcode"
                return 1
            fi
            echo "-H ${PKCS11_WRAPPER_SCRIPT} -p ${CUSTOMER_PUBLIC_KEY_FILE}"
            ;;
        pem-wrapped)
            # Wrapped key: sign via the helper (HSM-style), supplying the
            # derived public key just like the PKCS#11 path.
            if [ -z "${CUSTOMER_PUBLIC_KEY_FILE}" ]; then
                log "ERROR: Public key not available for rpi-sign-bootcode"
                return 1
            fi
            echo "-H ${PEM_WRAPPER_SCRIPT} -p ${CUSTOMER_PUBLIC_KEY_FILE}"
            ;;
        pem)
            echo "-k ${CUSTOMER_KEY_FILE_PEM}"
            ;;
        *)
            mark_permanent_failure
            log "ERROR: No signing key configured (bootcode signing requires a key)"
            return 1
            ;;
    esac
}

# Legacy compatibility wrapper - calls get_openssl_sign_args
# Deprecated: Use get_openssl_sign_args() directly
# This function exists for backward compatibility during migration
get_signing_directives() {
    if [ "${SIGNING_MODE}" = "" ]; then
        # Signing context not initialized - use legacy behavior
        if [ -n "${CUSTOMER_KEY_PKCS11_NAME}" ]; then
            echo "${CUSTOMER_KEY_PKCS11_NAME} -provider pkcs11 -provider default"
        elif [ -n "${CUSTOMER_KEY_FILE_PEM}" ]; then
            if [ -f "${CUSTOMER_KEY_FILE_PEM}" ]; then
                echo "${CUSTOMER_KEY_FILE_PEM} -keyform PEM"
            else
                log "ERROR: RSA private key \"${CUSTOMER_KEY_FILE_PEM}\" not a file."
                return 1
            fi
        else
            log "ERROR: Neither PKCS11 key name, or PEM key file specified."
            return 1
        fi
    else
        get_openssl_sign_args
    fi
}

# Get the kernel modules list file path
# Returns the path to the kernel modules list file, checking user config first
get_kernel_modules_list() {
    if [ -f /etc/rpi-sb-provisioner/kernel_modules.list ]; then
        echo "/etc/rpi-sb-provisioner/kernel_modules.list"
    else
        echo "/var/lib/rpi-sb-provisioner/kernel_modules.list"
    fi
}

# Find the kernel modules directory within a rootfs
# Checks both /lib/modules and /usr/lib/modules (older vs newer layouts)
#
# Arguments:
#   $1 - Root filesystem base directory
#
# Returns:
#   Prints the relative path to modules directory (e.g., "lib/modules" or "usr/lib/modules")
#   Returns 1 if no modules directory found
find_modules_dir() {
    _base="$1"
    
    # Check usr/lib/modules first (newer layout, e.g., Debian Bookworm+)
    if [ -d "${_base}/usr/lib/modules" ]; then
        echo "usr/lib/modules"
        return 0
    fi
    
    # Check lib/modules (older layout, e.g., Debian Bullseye and earlier)
    if [ -d "${_base}/lib/modules" ]; then
        echo "lib/modules"
        return 0
    fi
    
    return 1
}

# Find kernel version from a rootfs modules directory
# Checks both /lib/modules and /usr/lib/modules
#
# Arguments:
#   $1 - Root filesystem base directory
#
# Returns:
#   Prints the kernel version string (e.g., "6.6.31+rpt-rpi-v8")
#   Returns 1 if no kernel version found
find_kernel_version() {
    _base="$1"
    _modules_dir=$(find_modules_dir "${_base}") || return 1
    
    # Find the first kernel version directory
    _version=$(find "${_base}/${_modules_dir}" -mindepth 1 -maxdepth 1 -type d -exec basename {} \; 2>/dev/null | head -1)

    if [ -n "${_version}" ]; then
        echo "${_version}"
        return 0
    fi

    return 1
}

# Find ALL kernel version directories in a rootfs modules directory
# Checks both /lib/modules and /usr/lib/modules
#
# A single Raspberry Pi OS image ships modules for more than one kernel
# flavour (e.g. ...-v8 for Pi 3/4 and ...-2712 for Pi 5/CM5). The firmware
# picks which kernel to load at boot based on the board, so we cannot know in
# advance which one a given target will run. The pre-boot-auth initramfs must
# therefore carry modules for every flavour the image provides; modprobe then
# selects the matching set at runtime via uname -r.
#
# Arguments:
#   $1 - Root filesystem base directory
#
# Returns:
#   Prints each kernel version string on its own line
#   Returns 1 if no kernel version found
find_all_kernel_versions() {
    _base="$1"
    _modules_dir=$(find_modules_dir "${_base}") || return 1

    _versions=$(find "${_base}/${_modules_dir}" -mindepth 1 -maxdepth 1 -type d -exec basename {} \; 2>/dev/null)

    if [ -n "${_versions}" ]; then
        echo "${_versions}"
        return 0
    fi

    return 1
}

# Copy kernel modules with automatic dependency resolution
# This function uses rpi-modcopy to copy kernel modules and their dependencies
# into the destination directory (typically an initramfs).
#
# Arguments:
#   $1 - Source modules base directory (e.g., /path/to/rootfs containing lib/modules or usr/lib/modules)
#   $2 - Destination directory (e.g., /path/to/initramfs)
#   $3 - Kernel version string (e.g., "6.6.31+rpt-rpi-v8")
#   $4 - (Optional) Path to kernel modules list file. If not provided, uses default.
#
# The modules list file should contain one module name per line (without .ko extension).
# Lines starting with # are treated as comments.
#
# Returns:
#   0 - Success
#   1 - Failure (missing required files or commands)
copy_kernel_modules_with_deps() {
    _src_basedir="$1"
    _dst_basedir="$2"
    _kernel_version="$3"
    _modules_list="${4:-$(get_kernel_modules_list)}"

    # Validate inputs
    if [ -z "${_src_basedir}" ] || [ -z "${_dst_basedir}" ] || [ -z "${_kernel_version}" ]; then
        log "ERROR: copy_kernel_modules_with_deps requires source dir, dest dir, and kernel version"
        return 1
    fi

    # Find where modules are located (lib/modules or usr/lib/modules)
    # Use || true to prevent set -e from killing the script before the error check below
    _modules_rel_dir=$(find_modules_dir "${_src_basedir}") || true
    if [ -z "${_modules_rel_dir}" ]; then
        log "ERROR: No modules directory found in ${_src_basedir} (checked lib/modules and usr/lib/modules)"
        return 1
    fi

    if [ ! -d "${_src_basedir}/${_modules_rel_dir}/${_kernel_version}" ]; then
        log "ERROR: Kernel modules directory not found: ${_src_basedir}/${_modules_rel_dir}/${_kernel_version}"
        return 1
    fi

    if [ ! -f "${_modules_list}" ]; then
        log "ERROR: Kernel modules list file not found: ${_modules_list}"
        return 1
    fi

    log "Copying kernel modules using rpi-modcopy for ${_kernel_version}..."

    # Check rpi-modcopy explicitly rather than leaning on set -e: callers
    # invoke this function in an `if` (best-effort per flavour), which disables
    # set -e for the call, so a swallowed failure would otherwise fall through
    # to the success return below.
    if ! rpi-modcopy \
        --kernel-version="${_kernel_version}" \
        --module-dir="${_modules_rel_dir}" \
        --module-file="${_modules_list}" \
        "${_src_basedir}" \
        "${_dst_basedir}"; then
        log "ERROR: rpi-modcopy failed for ${_kernel_version}"
        return 1
    fi

    log "Kernel modules copied successfully"
    return 0
}

# Ensure `lock_device_private_key=1` is present in the given config.txt
# (which must live inside the signed boot.img the firmware authenticates).
# Firmware reads this key before any userspace runs and sets the OTP
# ECDSA key-slot to LOCKED, which blocks the raw-key export API for the
# rest of the boot session. Required by the rpi-verity-verifier boot-time
# runtime check.
#
# Modes:
#   enforce  — append the line under [all] if missing; no-op if present.
#   warn     — log a warning if missing; do not modify.
#
# Parameters:
#   $1 - path to config.txt (inside the boot.img mount)
#   $2 - mode: "enforce" or "warn" (default: "enforce")
ensure_lock_device_private_key() {
    _cfg="$1"
    _mode="${2:-enforce}"

    if [ ! -f "${_cfg}" ]; then
        log "ensure_lock_device_private_key: ${_cfg} not found; skipping"
        return 1
    fi

    # Match `lock_device_private_key=1` anywhere in the file (case-sensitive,
    # allow trailing whitespace/comments). Lines like "lock_device_private_key=0"
    # or any other value are NOT considered present — we want exactly =1.
    if grep -Eq '^[[:space:]]*lock_device_private_key[[:space:]]*=[[:space:]]*1([[:space:]]|#|$)' "${_cfg}"; then
        log "lock_device_private_key=1 already present in ${_cfg}"
        return 0
    fi

    if [ "${_mode}" = "warn" ]; then
        log "WARNING: lock_device_private_key=1 is MISSING from ${_cfg}."
        log "  If rpi-verity-verifier is installed in the image, boot will abort with KEY_NOT_LOCKED."
        log "  Add the line under [all] in your image's config.txt when using verity verification."
        return 0
    fi

    # Defensive: refuse to enforce if a conflicting value is already present.
    if grep -Eq '^[[:space:]]*lock_device_private_key[[:space:]]*=[[:space:]]*[^1][[:space:]]*' "${_cfg}"; then
        log "ERROR: ${_cfg} contains a conflicting lock_device_private_key setting."
        log "  Refusing to overwrite; correct the image's config.txt manually."
        return 1
    fi

    log "Appending lock_device_private_key=1 to ${_cfg} (under [all])"
    # Ensure there's an [all] section; if not, add one at the end.
    if ! grep -Eq '^[[:space:]]*\[all\][[:space:]]*$' "${_cfg}"; then
        printf '\n[all]\n' >> "${_cfg}"
    fi
    printf 'lock_device_private_key=1\n' >> "${_cfg}"
    return 0
}

# =============================================================================
# Signed boot.img preparation for IDP slot boot partitions
# =============================================================================
# IDP artefacts ship raw VFAT images for boot slots (containing firmware files,
# kernel, dtbs, config.txt, cmdline.txt). A signed-boot Pi 5 EEPROM rejects
# those and only loads a signed boot.img + boot.sig pair from the boot slot.
#
# This helper takes a source sparse VFAT, bundles its contents into a single
# boot.img via rpi-make-boot-image, signs it with the configured customer key,
# and emits a new sparse VFAT (sized to match the source) containing only
# {boot.img, boot.sig} suitable for direct flashing in place of the original.
#
# Caches output by content hash of source + key fingerprint so that, for an
# unchanged artefact and key, repeated runs reuse the prior signed copy.
#
# Arguments:
#   $1 - path to source sparse VFAT (e.g. ${IDP_DIR}/boot.sparse)
#   $2 - path where the signed sparse should land (caller-owned)
#
# Requires (caller responsibility): init_signing_context already invoked,
# RPI_DEVICE_FAMILY set, simg2img/img2simg/mkfs.fat/rpi-make-boot-image
# available on PATH.
prepare_signed_boot_simg() {
    _src_sparse="$1"
    _out_sparse="$2"

    if ! signing_available; then
        log "ERROR: prepare_signed_boot_simg requires a configured signing key"
        return 1
    fi

    if [ ! -f "${_src_sparse}" ]; then
        log "ERROR: source sparse not found: ${_src_sparse}"
        return 1
    fi

    if [ -f "${_out_sparse}" ]; then
        return 0
    fi

    # POSIX errexit is suspended inside a function called via "if !", which is
    # how callers invoke us. Check every critical step explicitly so a failure
    # cannot be silently turned into a 0-byte sparse cache entry.
    _fail() {
        log "ERROR: prepare_signed_boot_simg: $1"
        rm -rf "${_work}" 2>/dev/null
        return 1
    }

    _work=$(make_temp_dir "rpi-sb-bootimg.XXX") || return 1
    _src_img="${_work}/source.vfat"
    _src_mnt="${_work}/src-mnt"
    _boot_img="${_work}/boot.img"
    _boot_sig="${_work}/boot.sig"
    _out_vfat="${_work}/output.vfat"
    _out_mnt="${_work}/out-mnt"
    mkdir -p "${_src_mnt}" "${_out_mnt}" || _fail "mkdir mountpoints" || return 1

    # Inflate the IDP sparse so we can mount and walk its files.
    simg2img "${_src_sparse}" "${_src_img}" || _fail "simg2img source" || return 1

    # Mount read-only; rpi-make-boot-image only reads.
    mount -o loop,ro -t vfat "${_src_img}" "${_src_mnt}" \
        || _fail "mount source vfat" || return 1

    # Bundle every file in the boot VFAT (firmware, kernel, dtbs, cmdline.txt,
    # config.txt) into a single FIT-style boot.img the EEPROM will accept.
    rpi-make-boot-image -b "pi${RPI_DEVICE_FAMILY}" -a 64 \
        -d "${_src_mnt}" -o "${_boot_img}" \
        || { umount "${_src_mnt}" 2>/dev/null; _fail "rpi-make-boot-image"; return 1; }

    # Sidecar signature in the same format rpi-sb-provisioner.sh produces.
    sha256sum "${_boot_img}" | awk '{print $1}' > "${_boot_sig}" \
        || { umount "${_src_mnt}" 2>/dev/null; _fail "sha256sum boot.img"; return 1; }
    printf 'rsa2048: ' >> "${_boot_sig}"
    sign_image_hex "${_boot_img}" >> "${_boot_sig}" \
        || { umount "${_src_mnt}" 2>/dev/null; _fail "sign boot.img"; return 1; }
    validate_sig_file "${_boot_sig}" "IDP boot.sig" \
        || { umount "${_src_mnt}" 2>/dev/null; _fail "validate boot.sig"; return 1; }

    umount "${_src_mnt}" || _fail "umount source vfat" || return 1

    # Build a fresh VFAT sized to the original partition so the on-device
    # filesystem fits the existing slot exactly. -s 1 (one sector per cluster)
    # is required for compatibility with how rpi-make-boot-image lays out the
    # inner FAT, but at boot-slot sizes (typ. 128MB) that cluster count exceeds
    # FAT16's 65525-cluster ceiling, so we must force FAT32 explicitly. Without
    # -F 32 mkfs.fat fails with "Not enough or too many clusters" and -- with
    # errexit suspended -- the function would otherwise continue and write a
    # 40-byte stub sparse to the cache, leaving the device with a partition
    # the EEPROM finds boot.img in but cannot actually read.
    _src_size_bytes=$(stat -c%s "${_src_img}") \
        || _fail "stat source size" || return 1
    truncate -s "${_src_size_bytes}" "${_out_vfat}" \
        || _fail "truncate output vfat" || return 1
    mkfs.fat -s 1 -F 32 -n "BOOT" "${_out_vfat}" >/dev/null \
        || _fail "mkfs.fat -s 1 -F 32 (size=${_src_size_bytes})" || return 1

    mount -o loop -t vfat "${_out_vfat}" "${_out_mnt}" \
        || _fail "mount output vfat" || return 1
    cp "${_boot_img}" "${_out_mnt}/boot.img" \
        || { umount "${_out_mnt}" 2>/dev/null; _fail "cp boot.img"; return 1; }
    cp "${_boot_sig}" "${_out_mnt}/boot.sig" \
        || { umount "${_out_mnt}" 2>/dev/null; _fail "cp boot.sig"; return 1; }
    # The Pi 5 EEPROM only chainloads boot.img when config.txt in the
    # boot slot has boot_ramdisk=1. Without it the bootloader looks for
    # the legacy start4.elf/kernel_2712.img layout, which we replaced.
    # Mirrors get_fastboot_config_file() in the per-script callers.
    if [ -f /etc/rpi-sb-provisioner/boot_ramdisk_config.txt ]; then
        cp /etc/rpi-sb-provisioner/boot_ramdisk_config.txt "${_out_mnt}/config.txt"
    elif [ -f /var/lib/rpi-sb-provisioner/boot_ramdisk_config.txt ]; then
        cp /var/lib/rpi-sb-provisioner/boot_ramdisk_config.txt "${_out_mnt}/config.txt"
    else
        printf 'boot_ramdisk=1\n' > "${_out_mnt}/config.txt"
    fi
    sync
    umount "${_out_mnt}" || _fail "umount output vfat" || return 1

    # Atomic rename so a partial write never appears as a valid cache hit.
    img2simg -s "${_out_vfat}" "${_out_sparse}.tmp" \
        || _fail "img2simg output" || return 1
    mv "${_out_sparse}.tmp" "${_out_sparse}" \
        || _fail "mv output into cache" || return 1

    rm -rf "${_work}"
    return 0
}

# ---------------------------------------------------------------------------
# USB port restriction
#
# On a fixed programming jig only a known set of USB ports is wired to a
# device head; anything else plugged into the host (a keyboard, a USB stick,
# an engineer's laptop cable) must never be picked up and programmed. The
# restriction is expressed as drop-in rule files rather than a config
# variable so that a jig-specific package can ship the topology for its
# hardware without touching /etc/rpi-sb-provisioner/config.
#
# Two directories are consulted, in ascending precedence:
#
#   /usr/share/rpi-sb-provisioner/usb-ports.d/   package-supplied rules
#   /etc/rpi-sb-provisioner/usb-ports.d/         local administrator rules
#
# Only files named *.conf are read. Files are merged by basename, with the
# /etc copy winning, so a local empty file of the same name masks a
# package-supplied one entirely -- the same convention systemd uses for its
# own drop-ins.
#
# Each line is a USB topology path in the form the kernel uses for a device
# directory under /sys/bus/usb/devices ("<bus>-<port>[.<port>...]", e.g.
# "1-1.2"). Shell glob metacharacters are permitted, so "1-1.*" covers every
# downstream port of the hub on bus 1 port 1. Blank lines and lines whose
# first non-whitespace character is '#' are ignored.
#
# If no rule file yields a single pattern -- the default for a fresh install,
# where neither directory contains a *.conf -- the restriction is inactive
# and every port is accepted, preserving the historical behaviour.
USB_PORT_RULES_VENDOR_DIR="${USB_PORT_RULES_VENDOR_DIR:-/usr/share/rpi-sb-provisioner/usb-ports.d}"
USB_PORT_RULES_LOCAL_DIR="${USB_PORT_RULES_LOCAL_DIR:-/etc/rpi-sb-provisioner/usb-ports.d}"

# State for the "port excluded" outcome, recorded against the stage that
# rejected the device so the WebUI can distinguish a deliberate skip from a
# provisioning failure.
export PORT_EXCLUDED="PORT-EXCLUDED"

# List the effective rule files, one absolute path per line, after applying
# the /etc-masks-/usr/share precedence: local files first, then any package
# file whose basename the local directory did not already claim. Ordering is
# deterministic because pathname expansion sorts each directory's matches,
# though it does not affect the result -- the patterns are a set.
#
# Output (stdout): newline-separated absolute paths, possibly empty.
# Exit:   always 0.
list_usb_port_rule_files() {
    _upr_seen=""
    for _upr_dir in "${USB_PORT_RULES_LOCAL_DIR}" "${USB_PORT_RULES_VENDOR_DIR}"; do
        [ -d "${_upr_dir}" ] || continue
        for _upr_file in "${_upr_dir}"/*.conf; do
            [ -f "${_upr_file}" ] || continue
            _upr_base=$(basename "${_upr_file}")
            # A basename already claimed by the higher-precedence directory
            # masks the lower one. The guard newlines keep the substring
            # test from matching a partial basename.
            case "${_upr_seen}" in
                *"
${_upr_base}
"*) continue ;;
            esac
            _upr_seen="${_upr_seen}
${_upr_base}
"
            printf '%s\n' "${_upr_file}"
        done
    done
    return 0
}

# Read every effective rule file and emit the patterns they contain, one per
# line, stripped of comments and surrounding whitespace.
#
# Output (stdout): newline-separated glob patterns, possibly empty.
# Exit:   always 0.
list_usb_port_patterns() {
    list_usb_port_rule_files | while IFS= read -r _upp_file; do
        # `|| [ -n ... ]` so a final line without a trailing newline is kept.
        while IFS= read -r _upp_line || [ -n "${_upp_line}" ]; do
            # Strip a trailing comment, then surrounding whitespace.
            _upp_line=${_upp_line%%#*}
            _upp_line=$(printf '%s' "${_upp_line}" | tr -d '[:space:]')
            [ -n "${_upp_line}" ] || continue
            printf '%s\n' "${_upp_line}"
        done < "${_upp_file}"
    done
    return 0
}

# Report whether a USB port restriction is in force.
#
# Exit:   0 if at least one pattern is configured, 1 if unrestricted.
usb_port_restriction_active() {
    [ -n "$(list_usb_port_patterns | head -n1)" ]
}

# Decide whether a device on a given USB topology path may be provisioned.
#
# Arguments:
#   $1 - USB topology path (e.g. "1-1.2"). May be empty if the caller could
#        not resolve one.
#
# Exit:   0 if the device is permitted, 1 if it must be skipped.
#
# An empty path is permitted while unrestricted, but rejected once rules are
# in force: an allowlist cannot be honoured for a device whose port is
# unknown, and on a jig the safe reading of "I cannot tell where this is
# plugged in" is "do not program it".
usb_port_permitted() {
    _upp_path="$1"

    usb_port_restriction_active || return 0

    if [ -z "${_upp_path}" ]; then
        return 1
    fi

    # The loop runs in a subshell, so a match cannot set a variable the
    # caller would see -- it signals through the subshell's exit status
    # instead. Draining the loop without a match must therefore exit 1
    # explicitly, since a while loop that ends because `read` hit EOF is
    # itself a success.
    list_usb_port_patterns | {
        while IFS= read -r _upp_pattern; do
            # shellcheck disable=SC2254 # deliberate glob match against the pattern
            case "${_upp_path}" in
                ${_upp_pattern}) exit 0 ;;
            esac
        done
        exit 1
    }
}

# Reject a device that is not on a permitted USB port: log the reason,
# record the skip against the calling stage, and return non-zero so the
# caller can exit without running its failure hooks.
#
# Arguments:
#   $1 - Stage name for the state record ("bootstrap" or "triage")
#   $2 - USB topology path, possibly empty
#   $3 - Device serial for the state record
#
# Exit:   always 1, so callers can write `check || exit 0`.
record_usb_port_exclusion() {
    _upe_stage="$1"
    _upe_path="$2"
    _upe_serial="$3"

    if [ -z "${_upe_path}" ]; then
        log "USB port restriction active but this device's port could not be determined; skipping ${_upe_stage}"
    else
        log "USB port ${_upe_path} is not listed in ${USB_PORT_RULES_LOCAL_DIR} or ${USB_PORT_RULES_VENDOR_DIR}; skipping ${_upe_stage}"
    fi
    log "Permitted USB ports: $(list_usb_port_patterns | tr '\n' ' ')"

    record_state "${_upe_serial}" "${PORT_EXCLUDED}" "${_upe_stage}"
    return 1
}
