#!/bin/bash

ME=${0##*/}

VERSION="26.04"

: "${PATH_LIST:=/usr/local/bin /usr/sbin /usr/bin /sbin /bin}"
: "${BIN_SUBDIR:=bin}"
: "${LIB_SUBDIR:=lib}"

: "${DEFAULT_PROG_LIST:=ntfs-3g eject kmod}"
: "${ENCRYPT_PROG_LIST:=cryptsetup}"
: "${ENCRYPT_PROG_LIBS:=libgcc_s.so}"

: "${DEFAULT_KMOD_LIST:=depmod insmod lsmod modinfo modprobe}"
: "${DEFAULT_KMOD_LIBS:=libkmod.so liblzma.so libz.so liblz4.so libzstd.so}"

: "${DEFAULT_XZ_LIST:=xz xzcat unxz}"
: "${DEFAULT_ZSTD_LIST:=zstd zstdcat unzstd}"

X86_64_DIR=/lib/x86_64-linux-gnu

VAR_LIST=""
VAR_LIST+="PATH_LIST BIN_SUBDIR LIB_SUBDIR "
VAR_LIST+="DEFAULT_PROG_LIST ENCRYPT_PROG_LIST DEFAULT_KMOD_LIST "
VAR_LIST+="ENCRYPT_PROG_LIBS DEFAULT_KMOD_LIBS "

#PRETEND=true


usage() {
    cat <<Usage
usage: $ME [options] <[programs]

Copy dynamic linked programs and their supporting libraries
into the /bin and /lib directories of the initrd.

Options:
    -C --clean         Remove existing libraries and programs
    -c --create        Create the --to directory if it doesn't exist
    -e --encrypt       Add cryptsetup to built-in list of programs
    -f --from=<dir>    get libs and programs from under <dir>
    -F --force         Skip missing programs
    -h --help          Show this usage
    -n --no-color      Turn off text colors
    -p --pretend       Don't actually copy anything
    -q --quiet         Suppress modprobe warnings
    -s --show          Show default values of environment variables
    -t --to=<dir>      Copy to /bin and /lib under <dir>
    -v --verbose       Show more information
    -V --version       Show version

A --to directory or a from --directory must be specified.
Short options can be stacked.  Example:

    $ME -cnp

Usage

    exit "${1:-0}"
}

show_version() {
    echo "$ME Version:" "$VERSION"
}

#------------------------------------------------------------------------------
# callback routine for evaluating command line args
#------------------------------------------------------------------------------
eval_argument() {
    local arg=$1 val=$2
        case $arg in
            -create|c) CREATE=true             ;;
             -clean|C) DO_CLEAN=true           ;;
           -encrypt|e) ADD_ENCRYPT=true        ;;
              -from|f) FROM_DIR=$val           ;;
             -from=*) FROM_DIR=$val            ;;
             -force|F) FORCE=true              ;;
              -help|h) usage                   ;;
          -no-color|n) NO_COLOR=true           ;;
           -pretend|p) PRETEND=true            ;;
             -quiet|q) QUIET=true              ;;
              -show|s) do_show        ; exit 0 ;;
                -to|t) TO_DIR=$val             ;;
               -to=*) TO_DIR=$val              ;;
           -verbose|v) VERBOSE=true            ;;
           -version|V) show_version   ; exit 0 ;;
                    *) fatal "Unknown parameter -$arg" ;;
    esac
}

#------------------------------------------------------------------------------
# callback routine for saying which cli args take a parameter
#------------------------------------------------------------------------------
takes_param() {
    case $1 in
        -from|f) return 0 ;;
          -to|t) return 0 ;;
    esac
    return 1
}

#------------------------------------------------------------------------------
#
#------------------------------------------------------------------------------
main() {
    [[ $# -eq 0 ]] && usage 0

    show_version

    local SHIFT SHORT_STACK="cCefFhnpqtvV"

    local SUDO
    [[ "$UID" = 0 ]] || SUDO=sudo

    read_params "$@"
    shift "$SHIFT"

    [[ -n "$NO_COLOR" ]] || set_colors

    local prog_list abs_prog_list
    case $# in
        0)  prog_list=$DEFAULT_PROG_LIST                          ;;
        *)  prog_list="$*"
            [[ "$DO_CLEAN" ]] || qsay "Installing %s program(s) from command line" $#
            ;;
    esac

    if [[ $# -eq 0 ]] && [[ -n "$ADD_ENCRYPT$DO_CLEAN" ]]; then
        vsay "Adding encryption programs to clean list"
        prog_list="$prog_list $ENCRYPT_PROG_LIST"
    fi

    if [[ $# -eq 0 ]] && [[ -n "$DO_CLEAN" ]]; then
        vsay "Adding kmod programs to clean list"
        prog_list="$prog_list $DEFAULT_KMOD_LIST"
    fi

    : "${TO_DIR:=.}"

    [[ "$CREATE" ]] && mkdir -p "$TO_DIR"
    test -e "$TO_DIR" || fatal "The --to directory %s does not exist.  Use --create to have it created" "$TO_DIR"

    local bin_dir=${TO_DIR%/}/$BIN_SUBDIR  lib_dir=${TO_DIR%/}/$LIB_SUBDIR

    if [[ "$DO_CLEAN" ]]; then
        qsay "Cleaning out all libraries and these program(s): %s" "$prog_list"
        do_clean "$TO_DIR" "$lib_dir" "$bin_dir" "$prog_list"
        exit 0
    fi

    test -d "$TO_DIR" || fatal "The --to directory %s parameter is not a directory." "$TO_DIR"

    # Find the absolute path to each program under FROM_DIR
    local prog  path  abs_prog  abs_prog_list missing_progs
    for prog in $prog_list; do
        local found=
        for path in $PATH_LIST; do
            abs_prog=$path/$prog
            test -x "${FROM_DIR%/}$abs_prog" || continue
            found=true
            abs_prog_list="$abs_prog_list $abs_prog"
            break
        done
        [[ -z "$found" ]] && missing_progs="$missing_progs${missing_progs:+ }$prog"
    done

    [[ -n "$missing_progs" ]] && warn_err "The following program(s) were not found: %s"  "$missing_progs"

    if [[ "$VERBOSE" ]]; then
        say "Full path to programs under %s:" "${FROM_DIR:-/}"
        for prog in $abs_prog_list; do
            say "    %s" "$prog"
        done
    fi

    my_mkdir "$bin_dir"
    my_mkdir "$lib_dir"

    # Copy each program to the bin directory ...
    # And find the libraries needed for each program
    for abs_prog in $abs_prog_list; do
        local from_prog basename_prog
        from_prog="${FROM_DIR%/}$abs_prog"
        basename_prog=$(basename "$abs_prog")

        qsay "  add program: %s" "$basename_prog"
        [[ -e "$bin_dir/$basename_prog" ]] && pretend rm "$bin_dir/$basename_prog"
        pretend cp "$from_prog" "$bin_dir/$basename_prog"
        # kmod handling
        if [[ "$basename_prog" = "kmod" ]] && [[ -x "$bin_dir/kmod" ]]; then
            local link
            for link in $DEFAULT_KMOD_LIST; do
                pretend ln -nsf kmod "$bin_dir/$link"
            done
        fi
        # xz handling
        if [[ "$basename_prog" = "xz" ]] && [[ -x "$bin_dir/xz" ]]; then
            local link
            for link in $DEFAULT_XZ_LIST; do
                pretend ln -nsf xz "$bin_dir/$link"
            done
        fi
        # zstd handling
        if [[ "$basename_prog" = "zstd" ]] && [[ -x "$bin_dir/zstd" ]]; then
            local link
            for link in $DEFAULT_ZSTD_LIST; do
                pretend ln -nsf zstd "$bin_dir/$link"
            done
        fi

        local raw_libs
        if [[ -n "$FROM_DIR" ]]; then
            local LINUX_ARG; LINUX_ARG=$(get_linux_arg "$from_prog")
            raw_libs=$($SUDO ${LINUX_ARG:+"$LINUX_ARG"} chroot "$FROM_DIR" /usr/bin/ldd "$abs_prog")
        else
            raw_libs=$(ldd "$abs_prog")
        fi

        local libs; libs=$(echo "$raw_libs" | sed -nr "s/.*=>\s*([^ ]+)\s.*/\1/p")

        if [[ -z "$libs" ]]; then
            warn_err "No libs found for %s" "$(basename "$abs_prog")"
            continue
        fi

        local lib
        # luks2 handling - needs libgcc_s.so.* ...
        if [[ "$basename_prog" = "cryptsetup" ]] && [[ -x "$bin_dir/cryptsetup" ]]; then
            vsay "luks2 handling"
            if echo "$ENCRYPT_PROG_LIBS" | grep -q "libgcc_s.so"; then
                find_extra_libs "$libs" "$ENCRYPT_PROG_LIBS" libs
            fi
        fi

        # kmod handling - maybe needs libkmod.so, liblzma.so, libzstd.so
        if [[ "$basename_prog" = "kmod" ]] && [[ -x "$bin_dir/kmod" ]]; then
            vsay "kmod handling"
            if echo "$DEFAULT_KMOD_LIBS" | grep -q "libkmod"; then
                find_extra_libs "$libs" "$DEFAULT_KMOD_LIBS" libs
            fi
        fi

        # For each library, copy the library and the symlink to the lib directory
        for lib in $libs; do
            vsay "copy_symlink  ${FROM_DIR:-/} $lib  $lib_dir"
            copy_symlink "${FROM_DIR:-/}" "$lib" "$lib_dir" "${LINUX_ARG:-}"
        done

        # copy in the ld-linux program specifed by the program
        local ld_linux; ld_linux=$(echo "$raw_libs" | sed -nr "s@^\s*(/lib[^ ]*/ld-linux[^ ]*) .*@\1@p")
        [[ -z "$ld_linux" ]] && continue

        local ld_dir; ld_dir="$TO_DIR"$(dirname "$ld_linux")
        vsay "ld_linux: %s" "$ld_linux"
        if [[ "${ld_dir##*/}" = "lib64" ]]; then
            my_mkdir "$ld_dir"
            pretend cp -a "${FROM_DIR%/}${ld_linux}" "$ld_dir"

            ld_dir=${ld_dir%/lib64}$X86_64_DIR
        fi
        vsay "ld_dir: %s" "$ld_dir"
        my_mkdir "$ld_dir"
        copy_symlink "$FROM_DIR" "$ld_linux" "$ld_dir" "${LINUX_ARG:-}"

    done
}

#------------------------------------------------------------------------------
# Find extra libraries matching a pattern and append them to a variable.
# Usage: find_extra_libs "$libs" "$lib_pattern_list" varname
#   $1 - current libs list (space-separated paths)
#   $2 - lib pattern list (e.g. "libgcc_s.so liblzma.so")
#   $3 - name of variable to append results to
#------------------------------------------------------------------------------
find_extra_libs() {
    local current_libs=$1 lib_patterns=$2
    local -n libs_ref=$3

    local lib libdir="" libname
    vsay "libs=$current_libs"
    for lib in $current_libs; do
        libname=${lib##*/}; libname=${libname%%.*}
        [[ -n "${libname##libc*}" ]] && continue  # get libdir from libc.so
        libdir=${lib%/*}
        break
    done
    vsay "libdir: %s" "$libdir"

    if [[ -n "$libdir" ]]; then
        local extra=() regex
        regex=$(echo "$lib_patterns" | \
                sed -r 's/[.]so//g; s:^\s*:.*/(:; s/\s*$/)[.]so.*/; s/\s+/|/g')
        readarray -t extra < <(
            cd "${FROM_DIR:-/}" 2>/dev/null && \
            find "${FROM_DIR:+.}$libdir" ! -type d \
            -regextype posix-extended \
            -regex "$regex" 2>/dev/null
        )
        vsay "libs+=%s" "${extra[*]#.}"
        libs_ref+=" ${extra[*]}"
    fi
}

#------------------------------------------------------------------------------
# Copy a symlink and what it points to into dest_dir.  This code seems too
# complicated for what it does.
#------------------------------------------------------------------------------
copy_symlink() {
    local from_dir=${1:-/}  symlink=$2  dest_dir=$3  linux_arg=${4:-}

    # This works for relative symlinks and absolute symlinks that are also
    # exist on the host
    local real_file; real_file=$($SUDO ${linux_arg:+"$linux_arg"} chroot "$from_dir" readlink -f "$symlink")

    # If it was not really a symlink then just copy and return
    if [[ -z "$real_file" ]] || [[ "$real_file" = "$symlink" ]]; then
        local dest_file; dest_file="$dest_dir"/"$(basename "$symlink")"
        [[ -e "$dest_file" ]] || pretend cp "$from_dir$symlink" "$dest_file" \
            || fatal "Failed to copy:\n  %s -->\n  %s" "$from_dir$symlink" "$dest_file"
        return
    fi

    # Otherwise, copy the file and then make a relative symlink to it
    local dest_file; dest_file="$dest_dir"/"$(basename "$real_file")"
    [[ -e "$dest_file" ]] || pretend cp -a "${from_dir%/}/${real_file#/}" "$dest_file" \
        || fatal "Failed to copy:\n  %s -->\n  %s" "$from_dir$real_file" "$dest_file"

    local sym_targ; sym_targ=$(basename "$real_file")
    local sym_sorc; sym_sorc="$dest_dir"/"$(basename "$symlink")"
    [[ -e "$sym_sorc" ]] || pretend ln -s "$sym_targ" "$sym_sorc" \
        || fatal "Failed to create symlink at:\n  %s" "$sym_sorc"
}

#------------------------------------------------------------------------------
# We may need to run "linux32" or "linux64" when we do the chroot to find
# dynamic libs.
#------------------------------------------------------------------------------
get_linux_arg() {
    local prog=$1

    local need_64 new_arch=linux32
    if file "$prog" | grep -q 'x86[_-]64'; then
        need_64=true
        new_arch=linux64
    fi
    local have_64
    uname -m | grep -q 'x86[_-]64' && have_64=true
    local linux_arg
    [[ "$need_64" != "$have_64" ]] && linux_arg=$new_arch
    echo "$linux_arg"
}

#------------------------------------------------------------------------------
# Try to remove all libs and binaries that were added in.  We actually try to
# remove ALL libs since mismatched libs will tend to waste space and cause
# trouble.
#------------------------------------------------------------------------------
do_clean() {
    local to_dir=$1  lib_dir=$2  bin_dir=$3  prog_list=$4

    local prog
    local restore=""

    if  echo " $prog_list " | grep -sq " kmod "; then
        vsay "Adding kmod symlinked programs"
        prog_list="$prog_list $DEFAULT_KMOD_LIST"
    fi

    if  echo " $prog_list " | grep -sq " xz "; then
        vsay "Adding xz symlinked programs"
        prog_list="$prog_list $DEFAULT_XZ_LIST"
    fi

    if  echo " $prog_list " | grep -sq " zstd "; then
        vsay "Adding zstd symlinked programs"
        prog_list="$prog_list $DEFAULT_ZSTD_LIST"
    fi

    prog_list=$(echo "$prog_list" | sed 's/\s/\n/g' | sort -u | tr  '\n' ' ' );
    prog_list=${prog_list% }

    local busybox="$bin_dir"/busybox
    for prog in $prog_list; do
        local full="$bin_dir/$prog"
        { test -e "$full" || test -L "$full"; }  && pretend rm -f "$full"
        if [[ -x "$busybox" ]] && { "$busybox" --list | grep -sq "^${prog}$"; }; then
            restore+="$prog "
        fi
    done
    restore="${restore% }"
    # restore busybox links
    [[ -n "$restore" ]] && qsay "Restoring these busybox symlink(s): %s" "$restore"
    local link
    for link in $restore; do
        local full="$bin_dir/$link"
        if [[ -x "$busybox" ]] && { "$busybox" --list | grep -sq "^${link}$"; }; then
            pretend ln -nsf busybox "$full"
        fi
    done

    local lib
    for lib in "$to_dir"/lib*/{lib,ld}*.so*  "$to_dir"/"$X86_64_DIR"/ld-*.so*; do
        { [[ -e "$lib" ]] || [[ -L "$lib" ]]; } && pretend rm -f "$lib"
    done

    local dir
    for dir in "$to_dir"/lib*/ "$bin_dir" "$to_dir$X86_64_DIR"; do
        test -d "$dir" || continue
        [[ -z "$(ls -A "$dir" 2>/dev/null)" ]] && pretend rmdir "$dir"
    done
}

#------------------------------------------------------------------------------
# Show default values of our environment variables
#------------------------------------------------------------------------------
do_show() {
    printf "%s environment variables\n" "$ME"
    local var val
    for var in $VAR_LIST; do
        #eval val=\$"$var"
        local -n ref="$var"; val="$ref"
        printf "%20s: %s\n" "$var" "$val"
    done
}

#------------------------------------------------------------------------------
# Only try to make a directory if it doesn't already exist.  This is slightly
# convenient and it helps make or verbose reporting cleaner.
#------------------------------------------------------------------------------
my_mkdir() {
    local dir=$1
    test -e "$dir" || pretend mkdir -p "$dir"
}

#------------------------------------------------------------------------------
# Only run the command if not PRETEND.  Show command if PRETEND or VERBOSE.
#------------------------------------------------------------------------------
pretend() {
    { [[ "$PRETEND" ]] || [[ "$VERBOSE" ]]; } && echo "$@"
    [[ "$PRETEND" ]] && return
    "$@"
}

#-------------------------------------------------------------------------------
# Send "$@".  Expects
#
#   SHORT_STACK               variable, list of single chars that stack
#   fatal(msg)                routine,  fatal("error message")
#   takes_param(arg)          routine,  true if arg takes a value
#   eval_argument(arg, [val]) routine,  do whatever you want with $arg and $val
#
# Sets "global" variable SHIFT to the number of arguments that have been read.
#-------------------------------------------------------------------------------
read_params() {
    # Most of this code is boiler-plate for parsing cmdline args
    SHIFT=0
    # These are the single-char options that can stack

    local arg val

    # Loop through the cmdline args
    while [[ $# -gt 0 ]] && [[ ${#1} -gt 0 ]] && [[ -z "${1##-*}" ]]; do
        arg=${1#-}
        shift
        SHIFT=$((SHIFT + 1))

        # Expand stacked single-char arguments
        case $arg in
            [$SHORT_STACK][$SHORT_STACK]*)
                if echo "$arg" | grep -q "^[$SHORT_STACK]\+$"; then
                    local old_cnt=$#
                    local ARGS=()
                    read -ra ARGS < <(echo "$arg" | sed -r 's/([a-zA-Z])/ -\1 /g')
                    set -- "${ARGS[@]}" "$@"
                    SHIFT=$((SHIFT - $# + old_cnt))
                    continue
                fi;;
        esac

        # Deal with all options that take a parameter
        if takes_param "$arg"; then
            [[ $# -lt 1 ]] && fatal "Expected a parameter after: -$arg"
            val=$1
            [[ -n "$val" ]] && [[ -z "${val##-*}" ]] \
                && fatal "Suspicious argument after -$arg: $val"
            SHIFT=$((SHIFT + 1))
            shift
        else
            case $arg in
                *=*)  val=${arg#*=} ;;
                  *)  val="???"     ;;
            esac
        fi

        eval_argument "$arg" "$val"
    done
}

say() {
    local fmt=$1  ; shift
    # ignore "Don't use variables in the printf format string." shellcheck warnings
    # shellcheck disable=SC2059
    printf "$fmt\n" "$@"
}

vsay() {
    [[ "$VERBOSE" ]] || return
    say "$@"
}

qsay() {
    [[ "$QUIET" ]] && return
    say "$@"
}

fatal() {
    local fmt=$1; shift
    say "$ME$err_co fatal error:$warn_co $fmt$nc_co" "$@" >&2
    exit 2
}

warn() {
    local fmt=$1; shift
    say "$ME$err_co warning:$warn_co $fmt$nc_co" "$@" >&2
}

warn_err() {
    if [[ "$FORCE" ]]; then
        warn "$@"
    else
        fatal "$@"
    fi
}

# ignore "unused variables ..." shellcheck warnings
# shellcheck disable=SC2034
set_colors() {
   local e; e=$(printf "\e")

         black="${e}[0;30m" ;    blue="${e}[0;34m" ;    green="${e}[0;32m" ;    cyan="${e}[0;36m" ;
           red="${e}[0;31m" ;  purple="${e}[0;35m" ;    brown="${e}[0;33m" ; lt_gray="${e}[0;37m" ;
       dk_gray="${e}[1;30m" ; lt_blue="${e}[1;34m" ; lt_green="${e}[1;32m" ; lt_cyan="${e}[1;36m" ;
        lt_red="${e}[1;31m" ; magenta="${e}[1;35m" ;   yellow="${e}[1;33m" ;   white="${e}[1;37m" ;
         nc_co="${e}[0m"    ;   brown="${e}[0;33m" ;

           m_co=$cyan
          hi_co=$white
          err_co=$red
         bold_co=$yellow
         warn_co=$yellow
}

main "$@"
