#!/usr/bin/env bash
# SPDX-License-Identifier: GPL-2.0-only
#
# lsinitcpio - dump the contents of an initramfs image
#

shopt -s extglob

_list='--list'
_optcolor=1
if (( MESON_TEST )); then
    _f_functions='./functions'
else
    _f_functions='@INITCPIO_LIBDIR@/functions'
fi

declare -A bsdcpio_options=(
    [list]='--list'
    [input]='-i'
    [quiet]='--quiet'
)

usage() {
    cat <<USAGE
lsinitcpio @VERSION@
usage: ${0##*/} [action] [options] <initramfs>

  Actions:
   -a, --analyze        analyze contents of image
   -c, --config         show configuration file image was built with
   -l, --list           list contents of the image (default)
   -x, --extract        extract image to disk

  Options:
   -h, --help           display this help
   -n, --nocolor        disable colorized output
   -V, --version        display version information
   -v, --verbose        more verbose output
       --cpio           list or extract only the main CPIO image
       --early          list or extract only the early CPIO image

USAGE
}

version() {
    cat <<EOF
lsinitcpio @VERSION@
EOF
}

decomp() {
    local image="$1" offset="$2"
    if (( offset > 0 )); then
        dd if="$image" skip=$((offset / 512)) 2>/dev/null | ${_compress:-cat} ${_compress:+-cd}
    else
        ${_compress:-cat} ${_compress:+-cd} "$image"
    fi
}

unpack_uki() {
    local image="$1" workdir="$2"

    command -v objcopy &>/dev/null || die "Unable to unpack UKI: missing 'objcopy'"

    objcopy \
        --dump-section .initrd="$workdir"/initcpio.img \
        --dump-section .cmdline="$workdir"/cmdline.txt \
        --dump-section .osrel="$workdir"/os-release.txt \
        --dump-section .uname="$workdir"/uname.txt \
        "$image" /dev/null 2>/dev/null
    printf '%s/initcpio.img' "$workdir"
}

# shellcheck source=functions
. "$_f_functions"

# override the die method from functions
die() {
    error "$@"
    exit 1
}

size_to_human() {
    awk -v size="$1" '
    BEGIN {
        suffix[1] = "B"
        suffix[2] = "KiB"
        suffix[3] = "MiB"
        suffix[4] = "GiB"
        suffix[5] = "TiB"
        count = 1

        while (size > 1024) {
            size /= 1024
            count++
        }

        sizestr = sprintf("%.2f", size)
        sub(/\.?0+$/, "", sizestr)
        printf("%s %s", sizestr, suffix[count])
    }'
}

detect_filetype() {
    local file="$1" offset="$2" bytes

    bytes="$(od -An -t x1 -N6 -j "$offset" "$file" | tr -dc '[:alnum:]')"

    case $bytes in
        '303730373031')
            # no compression
            echo
            return
            ;;
    esac

    compression="$(detect_compression "$file" "$offset")"

    if [[ -n "$compression" ]]; then
        echo "$compression"
        return
    fi

    # still nothing? hrmm, maybe the user goofed and it's a kernel
    if kver "$1" >/dev/null; then
        die "'%s' is a kernel image, not an initramfs image!" "$file"
    fi

    # out of ideas, we're done here.
    die "Unexpected file type. Are you sure '%s' is an initramfs?" "$file"
}

detect_uki() {
    local file="$1" bytes

    IFS='' read -r -n 2 -d '' bytes < "$file"

    if [[ "$bytes" = "MZ" ]]; then
        return 0
    fi
    return 1
}

detect_early_img() {
    local -i early=0
    # early cpio archives can have a file called early_cpio in them
    early="$(LC_ALL=C.UTF-8 bsdtar -xOf "$1" early_cpio 2>/dev/null)"
    if [[ -z "$early" ]]; then
        # but not all do
        if LC_ALL=C.UTF-8 bsdtar -tf "$1" 'kernel/*/microcode/*.bin' &>/dev/null; then
            early=1
        fi
    fi
    printf '%d' "$early"
}

skip_early_img() {
    local -i offset
    local bytes

    # find last entry of the early cpio archive, marked by an entry with the name TRAILER!!!
    offset="$(grep -Fabom1 'TRAILER!!!' "$1" | cut -d: -f1)" || return 1

    # skip past the TRAILER!!! (10 chars) and to the next multiple of 512 bytes
    # we can do this in blocks of 512 because bsdtar/bsdcpio use blocksizes that are a multiple of 512
    offset="$(awk -v offset="$offset" '
    function ceil(v) {
        return (v == int(v)) ? v : int(v)+1
    }
    BEGIN {
        offset += 10
        printf "%d", ceil(offset / 512) * 512
    }')"


    # read a byte, check if it's NUL
    # if it's not, we've made it through the NULs
    # otherwise, skip ahead to the start of the next block and repeat
    bytes="$(od -An -t x1 -N1 -j "$offset" "$1" | tr -dc '[:alnum:]')" || return 1
    until [[ "$bytes" != '00' ]]; do
        offset=$(( offset + 512 ))
        bytes="$(od -An -t x1 -N1 -j "$offset" "$1" | tr -dc '[:alnum:]')" || return 1
    done
    printf '%d' "$offset"
}

analyze_image() {
    local -a binaries explicitmod modules imagename
    local kernver ratio columns image="$1" optimage="$2" offset="$3" workdir="$4"
    columns="$(tput cols)"

    # fallback in case tput failed us
    columns="${columns:-80}"

    # instead of reading the size from the inode, insist on reading the entire
    # image to ensure that it's in the cache when we decompress. This avoids
    # variance in timing and makes the time spent reading from storage roughly
    # constant.
    zsize="$(dd if="$image" skip=$((offset / 512)) 2>/dev/null | wc -c)"

    # calculate compression ratio
    if [[ -n "$_compress" ]]; then
        decomptime="$(TIMEFORMAT=%R; { time decomp "$image" "$offset" >/dev/null; } 2>&1)"
        fullsize="$(decomp "$image" "$offset" | wc -c)"
        ratio=".$(( zsize * 1000 / fullsize % 1000 ))"
    fi

    # decompress the image since we need to read from it multiple times.
    install -dm755 "$workdir/initcpio"
    if (( offset > 0 )); then
        bsdtar -C "$workdir/initcpio" -xf "$image" || die 'Failed to decompress early cpio'
    fi
    decomp "$image" "$offset" | bsdtar -C "$workdir/initcpio" -xf - || die 'Failed to decompress image'


    shopt -s globstar nullglob
    modules=("$workdir/initcpio/usr/lib/modules"/*/{kernel/**,extramodules}/*.ko*)
    IFS=/ read -r _ _ _ kernver _ <<<"${modules[0]#$workdir/initcpio/}"
    shopt -u globstar nullglob

    mapfile -t binaries < <(LC_ALL=C.UTF-8 find "$workdir/initcpio/usr/bin" -type f -printf %f\\n)

    read -r version <"$workdir/initcpio/VERSION"

    # shellcheck disable=SC1090,SC1091
    . "$workdir/initcpio/config"

    # MODULES in /config is a string
    # shellcheck disable=SC2128
    read -ra explicitmod <<<"$MODULES"

    # print results
    imagename=("$optimage")
    [[ -L "$optimage" ]] && imagename+=(" -> $(readlink -e "$optimage")")
    msg 'Image: %s%s' "${imagename[@]}"
    [[ -n "$version" ]] && msg 'Created with mkinitcpio %s' "$version"
    msg 'Kernel: %s' "${kernver:-unknown}"
    [[ -f "$workdir"/uname.txt ]] && msg 'uname: %s' "$(tr -d '\0' < "$workdir"/uname.txt)"
    if [[ -f "$workdir"/os-release.txt ]]; then
        (
            unset -v PRETTY_NAME NAME ID VERSION_ID BUILD_ID
            # shellcheck disable=SC1091
            . "$workdir"/os-release.txt
            msg 'OS Release: %s (%s %s)' "${PRETTY_NAME:-${NAME:-linux}}" \
                "${ID:-linux}" "${VERSION_ID:-${BUILD_ID:-rolling}}"
        )
    fi
    [[ -f "$workdir"/cmdline.txt ]] && msg 'Command line: %s' "$(tr -d '\0' < "$workdir"/cmdline.txt)"
    if (( offset > 0 )); then
        # the zero-indexed offset of the main image is the same as the size of the early image
        msg 'Early CPIO: %s' "$(size_to_human "$offset")"
    else
        msg 'Early CPIO: none'
    fi
    msg 'Size: %s' "$(size_to_human "$zsize")"

    if [[ -n "$_compress" ]]; then
        msg 'Compressed with: %s' "$_compress"
        msg2 'Uncompressed size: %s (%s ratio)' "$(size_to_human "$fullsize")" "$ratio"
        msg2 'Estimated decompression time: %ss' "$decomptime"
    fi
    printf '\n'

    if (( ${#modules[@]} )); then
        modules=("${modules[@]##*/}")
        modules=("${modules[@]%.ko*}")
        msg 'Included modules (%d):' "${#modules[@]}"
        for mod in "${modules[@]}"; do
            printf '  %s' "$mod"
            in_array "${mod//_/-}" "${explicitmod[@]//_/-}" && printf ' [explicit]'
            printf '\n'
        done | LC_ALL=C.UTF-8 sort | column -c"${columns}"
        printf '\n'
    fi

    msg 'Included binaries (%d):' "${#binaries[@]}"
    printf '  %s\n' "${binaries[@]}" | LC_ALL=C.UTF-8 sort | column -c"${columns}"
    printf '\n'

    if [[ -n "$EARLYHOOKS" ]]; then
        msg 'Early hook run order:'
        printf '  %s\n' "$EARLYHOOKS"
        printf '\n'
    fi

    # shellcheck disable=SC2128
    if [[ -n "$HOOKS" ]]; then
        msg 'Hook run order:'
        printf '  %s\n' "$HOOKS"
        printf '\n'
    fi

    if [[ -n "$LATEHOOKS" ]]; then
        msg 'Late hook run order:'
        printf '  %s\n' "$LATEHOOKS"
        printf '\n'
    fi

    if [[ -n "$CLEANUPHOOKS" ]]; then
        msg 'Cleanup hook run order:'
        printf '  %s\n' "$CLEANUPHOOKS"
        printf '\n'
    fi

    if [[ -n "$EMERGENCYHOOKS" ]]; then
        msg 'Emergency hook run order:'
        printf '  %s\n' "$EMERGENCYHOOKS"
        printf '\n'
    fi
}

_opt_short='achlnVvx'
_opt_long=('analyze' 'config' 'cpio' 'early' 'help' 'list' 'nocolor' 'version' 'verbose' 'extract')

parseopts "$_opt_short" "${_opt_long[@]}" -- "$@" || exit
set -- "${OPTRET[@]}"
unset _opt_short _opt_long OPTRET

while :; do
    case $1 in
        -a | --analyze)
            _optanalyze=1
            ;;
        -c | --config)
            _optshowconfig=1
            ;;
        --cpio)
            _optcpio=1
            ;;
        --early)
            _optearly=1
            ;;
        -h | --help)
            usage
            exit 0
            ;;
        -l | --list)
            _optlistcontents=1
            ;;
        -n | --nocolor)
            _optcolor=0
            ;;
        -V | --version)
            version
            exit 0
            ;;
        -v | --verbose)
            bsdcpio_options['verbose']='--verbose'
            ;;
        -x | --extract)
            unset 'bsdcpio_options[list]'
            ;;
        --)
            shift
            break 2
            ;;
    esac
    shift
done

_optimage="$1"
_image="$_optimage"

if [[ -t 1 ]] && (( _optcolor )); then
    try_enable_color
fi

[[ -n "$_image" ]] || die 'No image specified (use -h for help)'
[[ -f "$_image" ]] || die "No such file: '%s'" "$_image"
[[ -r "$_image" ]] || die "Unable to read file: '%s'" "$_image"

case $(( _optanalyze + _optlistcontents + _optshowconfig )) in
    0)
        # default action when none specified
        _optlistcontents=1
        ;;
    [!1])
        die "Only one action may be specified at a time"
        ;;
esac

case $(( _optearly + _optcpio )) in
    0)
        # default to listing/extracting both images
        _optearly=1
        _optcpio=1
        ;;
    # if == 1, only early or main cpio will be extracted/listed
    # if >1, same as 0
esac

_workdir="$(mktemp -d --tmpdir lsinitcpio.XXXXXX)"
trap 'rm -rf "$_workdir"' EXIT

if detect_uki "$_image"; then
    _image="$(unpack_uki "$_image" "$_workdir")" || die "Failed to unpack file as UKI: '%s'" "$_optimage"
    [[ -f "$_image" ]] || die "File does not contain initcpio image: '%s'" "$_optimage"
fi

if (( $(detect_early_img "$_image") )); then
    _offset="$(skip_early_img "$_image")"
else
    _offset=0
    _optearly=0
fi

# read compression type of the cpio image
_compress="$(detect_filetype "$_image" "$_offset")" || exit

if (( _optanalyze )); then
    analyze_image "$_image" "$_optimage" "$_offset" "$_workdir"
elif (( _optshowconfig )); then
    decomp "$_image" "$_offset" | bsdtar xOf - buildconfig 2>/dev/null ||
        die 'Failed to extract config from image (mkinitcpio too old?)'
else
    if (( _optearly )); then
        bsdcpio "${bsdcpio_options[@]}" <"$_image"
    fi
    if (( _optcpio )); then
        decomp "$_image" "$_offset" | bsdcpio "${bsdcpio_options[@]}"
    fi
fi

# vim: set ft=sh ts=4 sw=4 et: