LLM-based fw stitcher

.gitignore

@@ -3,6 +3,7 @@
 *.tar.gz
 *.new
 tests/tmp_downloads/
+__pycache__/
 
 
 # Added by cargo

Dockerfile

@@ -84,9 +84,11 @@ RUN pip install --upgrade pip && \
       git+https://github.com/marin-m/vmlinux-to-elf \
       jefferson \
       gnupg \
+      openai \
       poetry \
       psycopg2-binary \
       pycryptodome \
+      pydantic \
       pylzma \
       pyyaml \
       setuptools \
@@ -164,13 +166,31 @@ RUN --mount=type=secret,id=github_token \
         "https://raw.githubusercontent.com/qkaiser/arpy/23faf88a88488c41fc4348ea2b70996803f84f40/arpy.py" \
         /usr/local/lib/python3.10/dist-packages/arpy.py
 
-# Copy wrapper script into container so we can copy out - note we don't put it on guest path
+# Copy host wrappers into container so we can copy them out via the install
+# scripts — not put on guest PATH (the in-container entry points are
+# fakeroot_fw2tar and fwstitch, which live in /usr/local/bin).
 COPY ./fw2tar /usr/local/src/fw2tar_wrapper
-# And add install helpers which generate shell commands to install it on host
-COPY ./src/resources/banner.sh ./src/resources/fw2tar_install ./src/resources/fw2tar_install.local /usr/local/bin/
+COPY ./fwstitch /usr/local/src/fwstitch_wrapper
+# Install helpers (emit shell scripts the user pipes to sh / sudo sh on host)
+COPY ./src/resources/banner.sh \
+     ./src/resources/fw2tar_install ./src/resources/fw2tar_install.local \
+     ./src/resources/fwstitch_install ./src/resources/fwstitch_install.local \
+     /usr/local/bin/
+RUN chmod +x /usr/local/bin/fwstitch_install /usr/local/bin/fwstitch_install.local
 # Warn on interactive shell sessions and provide instructions for install
 RUN echo '[ ! -z "$TERM" ] && [ -z "$NOBANNER" ] && /usr/local/bin/banner.sh' >> /etc/bash.bashrc
 
 COPY src/fakeroot_fw2tar /usr/local/bin/fakeroot_fw2tar
 
+# Multi-filesystem stitcher (utils/stitch). The package goes to /opt/fw2tar so
+# `python3 -m stitch ...` works; `fwstitch` is a thin shim on PATH. The
+# Python entry point auto-re-execs under fakeroot for the `shard`/`all`
+# subcommands so firmware uid/gid metadata is preserved.
+COPY ./utils/stitch /opt/fw2tar/stitch
+RUN printf '%s\n' \
+        '#!/bin/bash' \
+        'exec env PYTHONPATH="/opt/fw2tar${PYTHONPATH:+:$PYTHONPATH}" python3 -m stitch "$@"' \
+        > /usr/local/bin/fwstitch && \
+    chmod +x /usr/local/bin/fwstitch
+
 CMD ["/usr/local/bin/banner.sh"]

fwstitch

@@ -0,0 +1,237 @@
+#!/bin/bash
+# Host wrapper for the fw2tar shard/plan/apply stitcher inside the rehosting/fw2tar container.
+#
+# - Auto-mounts file and directory arguments under /host_<basename>
+# - Passes through LLM_BASE_URL / LLM_API_KEY / LLM_MODEL / LLM_INSECURE so the
+#   container's `plan` step can reach your local model server
+# - Uses --network host for plan/all commands so http://localhost:8000/v1 etc.
+#   on the host is reachable from inside the container
+# - Mirrors ./fw2tar style. Subcommands: shard | plan | apply | all
+#
+# Usage examples:
+#   ./fwstitch shard ./firmware.bin -o ./shards
+#   ./fwstitch plan ./shards
+#   ./fwstitch apply ./shards ./shards/stitch_plan.yaml --out ./fw.stitched.rootfs.tar.gz
+#   LLM_BASE_URL=http://localhost:8000/v1 LLM_MODEL=gpt-oss-120b \
+#     ./fwstitch all ./firmware.bin --shard-dir ./shards --out ./fw.stitched.rootfs.tar.gz
+set -eu
+
+image="rehosting/fw2tar"
+verbose=false
+network_host=false
+extra_docker_args=()
+env_file=""
+
+# Load a KEY=VALUE-style .env file. Process env wins over file contents — we
+# only set keys that aren't already in the environment.
+load_env_file() {
+    local f="$1"
+    [[ -f "$f" ]] || return 0
+    while IFS= read -r line || [[ -n "$line" ]]; do
+        # Strip leading whitespace
+        line="${line#"${line%%[![:space:]]*}"}"
+        [[ -z "$line" ]] && continue
+        [[ "${line:0:1}" == "#" ]] && continue
+        # Allow optional `export ` prefix
+        line="${line#export }"
+        if [[ "$line" =~ ^([A-Za-z_][A-Za-z0-9_]*)=(.*)$ ]]; then
+            local k="${BASH_REMATCH[1]}"
+            local v="${BASH_REMATCH[2]}"
+            # Strip surrounding single or double quotes (but only if matched)
+            if [[ ${#v} -ge 2 && "${v:0:1}" == "${v: -1}" && ( "${v:0:1}" == '"' || "${v:0:1}" == "'" ) ]]; then
+                v="${v:1:${#v}-2}"
+            fi
+            if [[ -z "${!k:-}" ]]; then
+                export "$k=$v"
+            fi
+        fi
+    done < "$f"
+}
+
+# Subcommands that need network access for the LLM server.
+needs_network() {
+    case "${1:-}" in
+        plan|all) return 0 ;;
+        *) return 1 ;;
+    esac
+}
+
+print_help() {
+    cat <<EOF
+Usage: fwstitch [WRAPPER FLAGS] <subcommand> [SUBCOMMAND FLAGS] ...
+Wrapper script for running fw2tar's stitcher in a Docker container.
+
+Wrapper flags (must precede the subcommand):
+  --image NAME    image to run (default: $image)
+  --network MODE  docker network mode (default: 'host' for plan/all, bridge otherwise)
+  --env-file PATH load env vars from a KEY=VALUE file (process env wins)
+  --verbose       print mappings + docker command
+  --wrapper-help  this message
+
+.env auto-discovery: ~/.config/fwstitch/.env then ./.env are loaded if present.
+The format is plain KEY=VALUE per line (#-comments and an optional 'export '
+prefix are allowed). Process env vars always win over file contents.
+
+Subcommands (forwarded to fwstitch inside the container):
+  shard FIRMWARE -o SHARD_DIR        extract per-fragment .tar.gz + manifest
+  plan SHARD_DIR                     drive an LLM to produce stitch_plan.yaml
+  apply SHARD_DIR PLAN_YAML --out X  build the unified stitched .tar.gz
+  all FIRMWARE --shard-dir D --out X end-to-end
+
+Env vars forwarded into the container for plan/all:
+  LLM_BASE_URL, LLM_API_KEY, LLM_MODEL, LLM_INSECURE
+
+Pass through any subcommand-level flags as usual; this wrapper does no
+parsing of them, it only auto-mounts file and directory arguments.
+EOF
+}
+
+# Parse wrapper-level flags. Stop at the first non-flag arg (the subcommand).
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --wrapper-help) print_help; exit 0 ;;
+        --image)        image="$2"; shift 2 ;;
+        --network)      extra_docker_args+=("--network" "$2"); network_host=true; shift 2 ;;
+        --env-file)     env_file="$2"; shift 2 ;;
+        --verbose)      verbose=true; shift ;;
+        --)             shift; break ;;
+        *)              break ;;
+    esac
+done
+
+# Auto-load .env. Each load_env_file only sets keys that aren't already in the
+# environment, so order = precedence: process env > --env-file > ./.env >
+# ~/.config/fwstitch/.env. We feed them most-specific-first so the more
+# specific source wins when it disagrees with a more generic one.
+if [[ -n "$env_file" ]]; then
+    load_env_file "$env_file"
+fi
+load_env_file "./.env"
+load_env_file "${HOME}/.config/fwstitch/.env"
+
+if [[ $# -eq 0 ]]; then
+    print_help
+    exit 1
+fi
+
+subcmd="$1"
+shift
+cmd=("$subcmd" "$@")
+
+# Auto-mount any arg that is an existing file or directory, plus the value
+# immediately after --out, --plan-out, --shard-dir, --from-extracted (which
+# may not exist yet but should be writable on the host).
+maps=()
+
+map_path() {
+    local arg="$1"
+    local abspath
+    if [[ -e "$arg" ]]; then
+        abspath=$(realpath "$arg")
+    elif [[ "$arg" = /* ]]; then
+        abspath="$arg"
+    else
+        abspath="$PWD/$arg"
+    fi
+    local host_dir
+    local guest_dir
+    if [[ -d "$abspath" ]]; then
+        host_dir="$abspath"
+        guest_dir="/host_$(basename "$abspath")"
+        # Replace with full guest path (the dir itself, since we mount it)
+        new_value="$guest_dir"
+    elif [[ -f "$abspath" ]]; then
+        host_dir=$(dirname "$abspath")
+        guest_dir="/host_$(basename "$host_dir")"
+        new_value="$guest_dir/$(basename "$abspath")"
+    else
+        # Path doesn't exist; assume it's an output. Create parent dir on host,
+        # then mount that.
+        local parent
+        parent=$(dirname "$abspath")
+        mkdir -p "$parent"
+        host_dir="$parent"
+        guest_dir="/host_$(basename "$parent")"
+        new_value="$guest_dir/$(basename "$abspath")"
+    fi
+    maps+=("$host_dir:$guest_dir")
+    REWRITTEN="$new_value"
+}
+
+# Rewrite cmd[] in place. We mount:
+#   * any cmd[i] that is an existing path
+#   * the value after --out / --plan-out / --shard-dir / --from-extracted
+out_flags=(--out --plan-out --shard-dir --from-extracted -o)
+for ((i=1; i<${#cmd[@]}; i++)); do
+    arg="${cmd[$i]}"
+    prev="${cmd[$((i-1))]}"
+    is_out_value=false
+    for f in "${out_flags[@]}"; do
+        if [[ "$prev" == "$f" ]]; then is_out_value=true; break; fi
+    done
+    if $is_out_value || [[ -e "$arg" ]]; then
+        if [[ "$arg" == -* ]]; then continue; fi
+        map_path "$arg"
+        cmd[$i]="$REWRITTEN"
+    fi
+done
+
+# Deduplicate and sort mappings (longest first to avoid shadowing).
+if [[ ${#maps[@]} -gt 0 ]]; then
+    IFS=$'\n' maps=($(printf '%s\n' "${maps[@]}" | sort -u))
+    unset IFS
+fi
+
+# Build docker command
+docker_cmd=(docker run --rm)
+docker_cmd+=(-u "$(id -u):$(id -g)")
+
+# Network: --network host for plan/all unless the user overrode it
+if ! $network_host && needs_network "$subcmd"; then
+    docker_cmd+=(--network host)
+fi
+docker_cmd+=("${extra_docker_args[@]}")
+
+# Mounts
+for m in "${maps[@]}"; do
+    docker_cmd+=(-v "$m")
+done
+
+# LLM env vars — accept both LLM_API_KEY and the shorter LLM_KEY.
+for v in LLM_BASE_URL LLM_API_KEY LLM_KEY LLM_MODEL LLM_INSECURE; do
+    if [[ -n "${!v-}" ]]; then
+        docker_cmd+=(-e "$v=${!v}")
+    fi
+done
+
+# fw2tar hash for traceability (matches ./fw2tar)
+hash=$(sha1sum "$0" | awk '{print $1}')
+docker_cmd+=(-e "FWSTITCH_HASH=$hash")
+
+docker_cmd+=("$image" fwstitch "${cmd[@]}")
+
+if $verbose; then
+    echo "Mappings:"
+    for m in "${maps[@]}"; do echo "  $m"; done
+    echo "Docker command:"
+    redacted=()
+    redact_next=false
+    for tok in "${docker_cmd[@]}"; do
+        if $redact_next; then
+            # Redact the VALUE of an env-var assignment that names a secret.
+            case "$tok" in
+                LLM_API_KEY=*|LLM_KEY=*|*TOKEN=*|*SECRET=*|*PASSWORD=*)
+                    redacted+=("${tok%%=*}=<redacted>") ;;
+                *) redacted+=("$tok") ;;
+            esac
+            redact_next=false
+        else
+            redacted+=("$tok")
+            [[ "$tok" == "-e" ]] && redact_next=true
+        fi
+    done
+    echo "  ${redacted[*]}"
+    echo
+fi
+
+exec "${docker_cmd[@]}"

src/resources/banner.sh

@@ -41,3 +41,8 @@ echo -e "  $ docker run rehosting/fw2tar fw2tar_install.local | sh\n"
 echo -e "${BOLD}${RED}${STARS}Step 2: Run ${GREEN}fw2tar${RESET}${STARS}"
 echo -e "  $ fw2tar --help\n"
 
+echo -e "${BOLD}${RED}${STARS}Optional: Install ${GREEN}fwstitch${RESET}${STARS} (LLM-driven multi-filesystem stitcher)"
+echo -e "  $ docker run rehosting/fw2tar fwstitch_install | sudo sh"
+echo -e "  $ docker run rehosting/fw2tar fwstitch_install.local | sh"
+echo -e "  $ fwstitch --wrapper-help\n"
+

src/resources/fwstitch_install

@@ -0,0 +1,13 @@
+#!/bin/bash
+# Emits a shell script that installs the fwstitch host wrapper system-wide.
+# Run as: docker run rehosting/fw2tar fwstitch_install | sudo sh
+# Use a unique heredoc delimiter — the wrapper itself contains an 'EOF' inside
+# its print_help() heredoc, which would otherwise terminate the outer heredoc
+# early.
+echo "#!/bin/bash"
+echo "cat << '__FWSTITCH_WRAPPER_EOF__' | sudo tee /usr/local/bin/fwstitch >/dev/null"
+printf "%s" "$(cat /usr/local/src/fwstitch_wrapper)"
+echo
+echo "__FWSTITCH_WRAPPER_EOF__"
+echo "sudo chmod +x /usr/local/bin/fwstitch"
+echo "echo \"fwstitch installed successfully to /usr/local/bin/fwstitch\""

src/resources/fwstitch_install.local

@@ -0,0 +1,23 @@
+#!/bin/bash
+# Emits a shell script that installs the fwstitch host wrapper to ~/.local/bin.
+# Run as: docker run rehosting/fw2tar fwstitch_install.local | sh
+echo "#!/bin/bash"
+echo "mkdir -p \$HOME/.local/bin"
+# Unique delimiter — the wrapper has an 'EOF' inside its own print_help heredoc.
+echo "cat << '__FWSTITCH_WRAPPER_EOF__' > \$HOME/.local/bin/fwstitch"
+printf "%s" "$(cat /usr/local/src/fwstitch_wrapper)"
+echo
+echo "__FWSTITCH_WRAPPER_EOF__"
+echo "chmod +x \$HOME/.local/bin/fwstitch"
+echo "case \":\$PATH:\" in"
+echo "  *:\"\$HOME/.local/bin\":*) ;;"
+echo "  *) echo 'export PATH=\"\$HOME/.local/bin:\$PATH\"' >> \$HOME/.bashrc ;;"
+echo "esac"
+
+echo 'BOLD=$(tput bold 2>/dev/null || true)'
+echo 'RESET=$(tput sgr0 2>/dev/null || true)'
+
+echo "echo \"\${BOLD}Success!\${RESET} fwstitch installed to ~/.local/bin/fwstitch. If ~/.local/bin wasn't on PATH yet, source ~/.bashrc (or open a new shell) to pick it up.\""
+echo "echo"
+echo "echo \"Try:\""
+echo "echo \"  \${BOLD}fwstitch --wrapper-help\${RESET}\""