#!/bin/bash

ME=${0##*/}
# shellcheck disable=SC2034  # kept for parity with sibling scripts
VERSION="26.04"
MP_ARGS=(-D -a)

usage() {
        cat <<Usage
Usage: $ME [options] [file]
Copies modules listed in file from /lib/modules/\$KERN to
./lib/modules/\$KERN in the CWD.  Also copies all dependencies.

If only one kernel directory exists then we choose that kernel
automatically.

If no --dir or -- option is given and there are no files on
the command line then "default mode" is used.  The list of
modules is from a default list and the kernel/drivers/ata
directory is included automatically.

Options:
    --                 Read modules list on stdin
    -c --count         Display a count of the modules found
    -d --dir=DIR       Grab all modules found under DIR
    -e --encrypt       Include extra modules for encrypted live-usb
    -f --from=FROM     copy from lib/modules under FROM/
    -h --help          Show this usage
    -k --kernel=KERN   Use KERN instead of \`uname -r\`
    -O --only-encrypt  Only add encryption modules
    -p --pretend       Don't actually copy anything
    -q --quiet         Suppress modprobe warnings
    -s --show          Show the default list
    -Q --very-quiet    Also suppress warnings about duplicate inputs
    -t --to=TODIR      Copy to TODIR instead of CWD
    -v --verbose       List modules as they are copied

Usage
    exit "${1:-0}"
}

DEFAULT_LIST="
battery
btrfs
cdrom
crc16
crc32c-generic
crc32-pclmul
crct10dif-common
crct10dif-generic
crct10dif-pclmul
ecb
ehci-hcd
ehci-pci
exfat
ext4
f2fs
fat
firewire-core
firewire-ohci
firewire-sbp2
fuse
hid
hid-apple
hid-belkin
hid-cherry
hid-generic
hid-hyperv
hid-logitech
hid-logitech-dj
hid-logitech-hidpp
hid-microsoft
hid-monterey
hid-samsung
hv_balloon
hv_vmbus
hv_utils
hv_sock
hv_storvsc
hv_netvsc
hyperv-keyboard
isofs
jbd2
jfs
libcrc32c
libahci
loop
mbcache
md-mod
msdos
nls-ascii
nls-cp437
nls-utf8
nvme
ohci-hcd
ohci-pci
overlay
pci-hyperv
pci-hyperv-intf
reiserfs
rtsx_pci
rtsx_pci_sdmmc
scsi-mod
scsi-transport-fc
sdhci
sd-mod
sg
squashfs
sr-mod
uas
udf
ufshcd-pci
uhci-hcd
ums-realtek
usb-common
usb-storage
usbcore
usbhid
usbmon
raid0
raid1
raid10
raid456
raid6_pq
raid_class
virtio-blk
virtio-pci
vfat
vmd
virtio-gpu
xcbc
xfs
xhci-hcd
xhci-pci
xor
"

CRYPT_LIST="
aes
async_memcpy
async_pq
async_raid6_recov
async_tx
async_xor
blowfish
dm-crypt
dm-mod
serpent
sha256
xts
"
# i915
# nouveau
# radeon

CONDITIONAL_MODULES="
hyperv_fb
linear
multipath
"

module_exists_in_tree() {
    local module_dir=$1
    local module_name=$2
    local file_name=${module_name//-/_}
    find "$module_dir" -type f -name "${file_name}.ko*" -print -quit | grep -q .
}

main() {

    local to_dir from_dir kernel

    local short_stack="defhkOpqQstv"
    while [[ $# -gt 0 ]] && [[ -n "$1" ]] && [[ -z "${1##-*}" ]]; do
        local arg=${1#-} val=
        shift

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

        case $arg in
            -dir|-from|-kernel|-to|[dfkt])
                [[ $# -lt 1 ]] && fatal "Expected a parameter after: -$arg"
                val=$1
                shift;;
            *=*)
                val=${arg#*=}
                arg=${arg%%=*} ;;
             *)
                 val="???" ;;
        esac

        case $arg in
                  -) USE_STDIN=true                  ;;
           -count|c) SHOW_COUNT=true                 ;;
             -dir|d) DIRS="${DIRS:+$DIRS,}$val"       ;;
         -encrypt|e) ADD_ENCRYPT=true                ;;
            -help|h) usage 0                         ;;
            -from|f) from_dir=${val%/}               ;;
          -kernel|k) kernel=$val                     ;;
    -only-encrypt|O) ONLY_ENCRYPT=true               ;;
         -pretend|p) PRETEND=true                    ;;
           -quiet|q) QUIET=true                      ;;
      -very-quiet|Q) VERY_QUIET=true                 ;;
            -show|s) echo "$DEFAULT_LIST" ; exit     ;;
              -to|t) to_dir=${val%/}                 ;;
         -verbose|v) VERBOSE=true                    ;;
                  *) fatal "Unknown argument -$arg"  ;;
        esac
    done

    { [[ "$QUIET" ]] || [[ "$VERY_QUIET" ]]; } && MP_ARGS+=(--quiet)
    # fill in default kernel if there is only one
    if [[ -z "$kernel" ]]; then
        local d="$from_dir/lib/modules"
        test -d "$d" || fatal "Could not find directory \"$d\""
        local kernels=()
        readarray -t kernels < <(find "$d" -mindepth 1 -maxdepth 1 -printf '%f\n')
        [[ ${#kernels[@]} -eq 1 ]] && kernel=${kernels[0]}
    fi

    [[ "$kernel" ]] && MP_ARGS+=(-S "$kernel")
    [[ "$from_dir" ]] && MP_ARGS+=(-d "$from_dir")

    : "${kernel:=$(uname -r)}"
    : "${to_dir:=.}"

    local mod_dir="$from_dir/lib/modules/$kernel"
    test -d "$mod_dir" || fatal "Directory $mod_dir does not exist"

    local list1 list2 default_mode
    # grab list from files or stdin and strip comments

    if [[ "$ADD_ENCRYPT" ]]; then
        DEFAULT_LIST=$DEFAULT_LIST$CRYPT_LIST
    elif [[ "$ONLY_ENCRYPT" ]]; then
        DEFAULT_LIST=$CRYPT_LIST
    fi

    local conditional_module
    for conditional_module in $CONDITIONAL_MODULES; do
        if module_exists_in_tree "$mod_dir" "${conditional_module//_/-}"; then
            DEFAULT_LIST="${DEFAULT_LIST}"$'\n'"$conditional_module"
        fi
    done

    if [[ -z "$USE_STDIN" ]] && [[ $# -eq 0 ]]; then
        default_mode=true
        list1=$(echo "$DEFAULT_LIST" | grep -v "^\s*#" | sed 's/\s*#.*//')
    else
        list1=$(cat "$@" | grep -v "^\s*#" | sed -e 's/\s*#.*//' -e "s/,/ /g")
    fi

    [[ "$default_mode" ]] && [[ -z "$DIRS" ]] && DIRS="kernel/drivers/ata,kernel/drivers/mmc"
    local d dir_array=()
    [[ -n "$DIRS" ]] && IFS=', ' read -ra dir_array <<< "$DIRS"
    for d in "${dir_array[@]}"; do
        local dir=$mod_dir/$d
        test -d "$dir" || fatal "Subdirectory $dir does not exist"
        list1="$list1
$(find "$dir" -type f -name '*\.ko*' -printf "%f\n" | sed -nr 's/\.ko(|\.[[:alpha:]]+)$//p')"
    done

    local unique_list
    unique_list=$(printf '%s\n' "$list1" | sed 's/_/-/g' | sort -u)

    list2=$(printf '%s\n' "$unique_list" \
        | xargs /sbin/modprobe -n "${MP_ARGS[@]}" | sed -n 's/^insmod //p' | cut -d' ' -f1 | sort -u)

    local full cnt=0
    [[ -z "${from_dir##/*}" ]] || from_dir="$PWD/$from_dir"
    while read -r full; do
        [[ "$VERBOSE" ]] && printf "  %s\n" "$(basename "$full")"
        local dir
        dir=$(dirname "$full")
        dir="$to_dir/${dir##"$from_dir"/}"
        test -d "$dir" || pretend mkdir -p "$dir"
        pretend cp "$full" "$dir/"
        cnt=$((++cnt))
    done<<List_2
$list2
List_2

    [[ "$SHOW_COUNT" ]] && printf '%s: %s module(s)\n' "$ME" "$cnt"

    exit 0
}

pretend() {
    [[ "$PRETEND" ]] && return
    "$@"
}

fatal() {
    local fmt=$1; shift
    # shellcheck disable=SC2059
    printf "$ME fatal error: $fmt\n" "$@" >&2
    exit 3
}

warn() {
    local fmt=$1; shift
    # shellcheck disable=SC2059
    printf "$ME warning: $fmt\n" "$@" >&2
}

main "$@"
