An automation bash script for chroot-ing into the rootfs partition of a given disk image

edit_rootfs.sh

#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
ARGPARSE_SCRIPT="$SCRIPT_DIR/argparse.bash"
if [[ ! -f "$ARGPARSE_SCRIPT" ]]; then
    >&2 echo "[ERROR] '$SCRIPT_DIR/argparse.bash' does not exist! Please download it from 'https://github.com/nhoffman/argparse-bash/raw/master/argparse.bash'."
    exit 1
fi

ARGPARSE_DESCRIPTION="Chroot into the ROOTFS_PARTITION of the given DISK_IMAGE."
source "$ARGPARSE_SCRIPT" || exit 1
argparse "$@" <<EOF || exit 1
parser.add_argument('disk_image', help='path of the disk image file, i.e., ./disk.img')
parser.add_argument('rootfs_partition', type=int, help='the partition number of the rootfs partition, i.e., 4')
parser.add_argument('-u', '--login-user', default='root', help='login to chroot shell as LOGIN_USER; the default value is "root"')
parser.add_argument('-a', '--fs-arch', default='native', help='the architecture of the binaries in rootfs, i.e., aarch64; the default value is "native"')
parser.add_argument('-z', '--zerofree', action='store_true', help='fill unused areas in the rootfs partition with zeroes')
EOF

if [[ ! -f "$DISK_IMAGE" ]]; then
    >&2 echo "[ERROR] The given disk image file '$DISK_IMAGE' does not exist!"
    exit 1
fi

dep_pkgs=""

if [[ "$ZEROFREE" = "yes" ]]; then
    dep_pkgs="$dep_pkgs zerofree"
    echo "Package 'zerofree' is required."
fi

require_qemu=0
qemu_cmd=""
local_arch="$(uname -m)"
if [[ "$FS_ARCH" = "native" || "$local_arch" = "$FS_ARCH" ]]; then
    require_qemu=0
    qemu_cmd=""
else
    dep_pkgs="$dep_pkgs qemu-user-static"
    require_qemu=1
    qemu_cmd="qemu-${FS_ARCH}-static"
    echo "Local architecture '$local_arch' is different from FS_ARCH '$FS_ARCH'. Package 'qemu-user-static' is required."
fi

echo "Checking if $dep_pkgs exist..."
if [[ $(dpkg -l dep_pkgs &> /dev/null) -ne 0 ]]; then
    echo "$dep_pkgs not found. Please install them with the package manager."
    exit 1
fi
echo "[OK]"

if [[ $require_qemu -eq 1 ]]; then
    echo "Checking if the command '${qemu_cmd}' exists..."
    command -v "${qemu_cmd}" || {
        >&2 echo "[ERROR] The command '${qemu_cmd}' does not exist while 'qemu-user-static' has been installed. Maybe a typo in FS_ARCH '${FS_ARCH}'?"
        exit 1
    }
fi

echo "Finding the first unused loop device..."
echo "sudo permission is required."
loop_dev="$(sudo losetup -f)"
echo "[OK]"
echo "Associating loop device '$loop_dev' with image file '$DISK_IMAGE'..."
sudo losetup --partscan "$loop_dev" "$DISK_IMAGE"
echo "[OK]"

# try
set +e; ( set -e;

    echo "Getting the partition info of '$loop_dev'..."
    lsblk "$loop_dev"
    partitions_num=$(($(lsblk -pno NAME,MOUNTPOINT "$loop_dev" | grep -c ' ')-1))
    if [[ "$partitions_num" -lt "$ROOTFS_PARTITION" ]]; then
        >&2 echo "[ERROR] The given disk image file '$DISK_IMAGE' must have at least '$ROOTFS_PARTITION' partition!"
        exit 1
    fi

    root_mntpnt=$(mktemp -d)
    echo "Mounting '/' (${loop_dev}p${ROOTFS_PARTITION}) partition to '$root_mntpnt'..."
    sudo mount "${loop_dev}p${ROOTFS_PARTITION}" "$root_mntpnt"
    sudo mount --bind /dev     "${root_mntpnt}/dev"
    sudo mount --bind /dev/pts "${root_mntpnt}/dev/pts"
    sudo mount --bind /sys     "${root_mntpnt}/sys"
    sudo mount --bind /proc    "${root_mntpnt}/proc"
    sudo mount --bind /run     "${root_mntpnt}/run"
    echo "[OK]"

    if [[ $require_qemu -eq 1 ]]; then
        copied_qemu=0
        qemu_path=/usr/bin/${qemu_cmd}
        if [[ ! -f "${root_mntpnt}/${qemu_path}" ]]; then
            sudo cp $qemu_path "${root_mntpnt}/${qemu_path}"
            copied_qemu=1
        fi
    fi

    # try
    set +e; ( set -e;

        cmd_after_chroot=" \
            set -e; \
            echo \"========== Welcome to the inner world! ==========\"; \
            echo \"========== Press 'Ctrl-D' or execute 'exit' to quit. ==========\"; \
            set +e; ( set -e; su -l ${LOGIN_USER} ); set -e; \
            echo \"========== Goodbye from the inner world! ==========\"; \
            "
        sudo chroot "${root_mntpnt}" ${qemu_cmd} /bin/bash -c "$cmd_after_chroot"

    # catch
    ); set -e;

    echo "Unmounting '/' (${loop_dev}p${ROOTFS_PARTITION}) partition from '$root_mntpnt'..."
    if [[ $require_qemu -eq 1 && $copied_qemu -eq 1 ]]; then
        sudo rm -f "${root_mntpnt}/${qemu_path}"
    fi
    sudo umount "${root_mntpnt}/dev/pts"
    sudo umount "${root_mntpnt}/dev"
    sudo umount "${root_mntpnt}/sys"
    sudo umount "${root_mntpnt}/proc"
    sudo umount "${root_mntpnt}/run"
    sudo umount "$root_mntpnt"
    rm -rf "$root_mntpnt"
    echo "[OK]"

    if [[ "$ZEROFREE" = "yes" ]]; then
        echo "Performing zerofree for '${loop_dev}p${ROOTFS_PARTITION}'..."
        sudo zerofree -vvv "${loop_dev}p${ROOTFS_PARTITION}"
    fi

# catch
); set -e;

echo "Detaching loop device '$loop_dev' from image file '$DISK_IMAGE'..."
sudo losetup -d "$loop_dev"
echo "[OK]"

echo "Script execution finished."