#!/bin/sh
#
# tiny initramfs
#

print()
{
    printf "%b %s\n" "${2:-"\033[1;37m>>\033[m"}" "$1"
}

panic()
{
    print "${1:-unexpected error occurred}" \
          "\033[1;31m!!\033[m"  >&2; exit 1
}

usage()
{
    cat << EOF
usage: ${0##*/} [option...]
       -o, --output <file> set initramfs output path
                           default is /boot/initramfs-$(uname -r)

       -c, --config <file> set config file path
                           default is /etc/tinyramfs/config

       -m, --modules <dir> set modules directory
                           default is /lib/modules

       -s, --sources <dir> set sources directory
                           default is /usr/share/tinyramfs

       -k, --kernel  <ver> set kernel version
                           default is $(uname -r)

       -H, --hooks   <dir> set hooks directory
                           default is /etc/tinyramfs/hooks (user hooks)
                           and /usr/share/tinyramfs/hooks (system hooks)

       -d, --debug         enable debug mode
       -f, --force         overwrite initramfs image

EOF
}

prepare_environment()
{
    while [ "$1" ]; do case "$1" in
        -o | --output)
            output="${2:?}"; shift 2
        ;;
        -c | --config)
            config="${2:?}"; shift 2
        ;;
        -m | --modules)
            moddir="${2:?}"; shift 2
        ;;
        -s | --sources)
            srcdir="${2:?}"; shift 2
        ;;
        -k | --kernel)
            kernel="${2:?}"; shift 2
        ;;
        -H | --hooks)
            hksdir="${2:?}"; shift 2
        ;;
        -d | --debug)
            debug=1; shift 1
        ;;
        -f | --force)
            force=1; shift 1
        ;;
        -h | --help)
            usage; exit 0
        ;;
        *)
            printf "invalid option: %s\n\n" "$1"

            usage; exit 1
        ;;
    esac; done

    print "preparing environment"

    . "${config:-/etc/tinyramfs/config}"

    : "${kernel:=$(uname -r)}"
    : "${moddir:=/lib/modules}"
    : "${srcdir:=/usr/share/tinyramfs}"
    : "${output:=/boot/tinyramfs-${kernel}}"

    mkdir -p "${tmpdir:=${TMPDIR:-/tmp}/tinyramfs.$$}"

    # false positive
    # shellcheck disable=2015
    [ "$debug" = 1 ] && set -x || trap "rm -rf $tmpdir" EXIT INT
}

prepare_initramfs()
{
    print "preparing initramfs"

    # make directories
    mkdir -p \
        "${tmpdir}/dev" \
        "${tmpdir}/sys" \
        "${tmpdir}/tmp" \
        "${tmpdir}/run" \
        "${tmpdir}/var" \
        "${tmpdir}/proc" \
        "${tmpdir}/root" \
        "${tmpdir}/usr/lib" \
        "${tmpdir}/usr/bin" \
        "${tmpdir}/mnt/root" \
        "${tmpdir}/etc/tinyramfs"

    # make symlinks
    ln -s usr/lib "${tmpdir}/lib"
    ln -s usr/bin "${tmpdir}/bin"
    ln -s usr/bin "${tmpdir}/sbin"
    ln -s ../run  "${tmpdir}/var/run"
    ln -s ../run/lock "${tmpdir}/var/lock"
    ln -s bin     "${tmpdir}/usr/sbin"

    # TODO make blkid optional
    # copy required binaries
    for _binary in \
        \[ sh ln env kill mkdir \
        blkid sleep mount printf \
        switch_root "${srcdir}/device-helper"
    do
        copy_binary "$_binary"
    done

    # copy init
    cp "${srcdir}/init" "${tmpdir}/init"
    chmod 755           "${tmpdir}/init"

    # copy config
    cp "$config" "${tmpdir}/etc/tinyramfs/config"
    chmod 600    "${tmpdir}/etc/tinyramfs/config"
}

copy_binary()
{
    binary=$(command -v "$1")

    # check if binary exist and builtin
    case "$binary" in */*) ;;
        "")
            panic "$1 does not exist"
        ;;
        *)
            # word splitting is safe by design
            # shellcheck disable=2086
            IFS=:; set -- $PATH; unset IFS

            # assume that `command -v` returned builtin command.
            # this behavior depends on shell implementation.
            # to be independented we simply iterating over PATH
            # to find external alternative ( e.g kill => /bin/kill )
            for _dir; do
                [ -x "${_dir}/${binary}" ] || ! continue

                binary="${_dir}/${binary}"; break
            done || panic "$1 does not exist"
        ;;
    esac

    # check if binary already exist
    [ -e "${tmpdir}/bin/${binary##*/}" ] && return 0

    # iterate throught symlinks and copy them
    while [ -h "$binary" ]; do symlink=$(ls -ld "$binary")
        cp -P "$binary" "${tmpdir}/bin" || panic
        binary="${binary%/*}/${symlink##* ->*[ /]}"
    done

    {
        cp "$binary" "${tmpdir}/bin"
        chmod 755    "${tmpdir}/bin/${binary##*/}"
    } || panic

    strip "${tmpdir}/bin/${binary##*/}" > /dev/null 2>&1 || :

    # copy binary dependencies if any
    ldd "$binary" 2> /dev/null |

    while read -r _library || [ "$_library" ]; do

        # skip unneeded stuff
        [ "${_library##*vdso*}" ] || continue

        _library="${_library#* => }"
        _library="${_library% *}"

        copy_library "$_library"
    done
}

copy_library()
{
    library="$1"

    # check if library already exist
    [ -e "${tmpdir}/lib/${library##*/}" ] && return 0

    # iterate throught symlinks and copy them
    while [ -h "$library" ]; do symlink=$(ls -ld "$library")
        cp -P "$library" "${tmpdir}/lib" || panic
        library="${library%/*}/${symlink##* ->*[ /]}"
    done

    {
        cp "$library" "${tmpdir}/lib"
        chmod 755     "${tmpdir}/lib/${library##*/}"
    } || panic

    strip "${tmpdir}/lib/${library##*/}" > /dev/null 2>&1 || :
}

copy_module()
{
    module="$1"

    modprobe -S "$kernel" -D "$module" 2> /dev/null |

    while read -r _ module || [ "$module" ]; do

        # check if module contains full path(not builtin)
        [ "${module##*/*}" ] && continue

        # check if module already exist
        [ -e "${tmpdir}${module}" ] && continue

        {
            mkdir -p     "${tmpdir}${module%/*}"
            cp "$module" "${tmpdir}${module}"
            chmod 644    "${tmpdir}${module}"
        } || panic
    done
}

copy_hook()
{
    hook="$1"

    for _dir in "$hksdir" /etc/tinyramfs/hooks /usr/share/tinyramfs/hooks; do
        [ -f "${_dir}/${hook}/${hook}" ] || ! continue

        print "running $hook hook"; . "${_dir}/${hook}/${hook}"

        if [ -f "${_dir}/${hook}/${hook}.init" ]; then
            mkdir -p "${tmpdir}/usr/share/tinyramfs/hooks/${hook##*/}"
            cp "${_dir}/${hook}/${hook}.init" \
               "${tmpdir}/usr/share/tinyramfs/hooks/${hook##*/}"

        elif [ -f "${_dir}/${hook}/${hook}.init.early" ]; then
            mkdir -p "${tmpdir}/usr/share/tinyramfs/hooks/${hook##*/}"
            cp "${_dir}/${hook}/${hook}.init.early" \
               "${tmpdir}/usr/share/tinyramfs/hooks/${hook##*/}"
        fi || panic

        break
    done || panic "could not run $hook"
}

copy_modules()
{
    # skip this function if kernel
    # compiled with builtin modules
    if [ "$monolith" = 1 ]; then
        print "skipping modules"
        return 0

    elif [ "$hostonly" = 1 ]; then
        print "copying hostonly modules"

        # perform autodetection of modules via /sys
        # see https://wiki.archlinux.org/index.php/Modalias
        find /sys/devices -name modalias -exec sort -u {} + |

        while read -r _module || [ "$_module" ]; do

            # skip unneeded modules and skip modules which
            # depends on them as well
            case $(modprobe -S "$kernel" -D "$_module") in
                *wmi* | *gpu* | *net*) continue ;;
            esac 2> /dev/null

            copy_module "$_module"
        done

        # copy root filesystem module
        if [ "$root_type" ]; then
            copy_module "$root_type"
        else
            while read -r _ _dir _type _; do
                [ "$_dir" = / ] || ! continue

                copy_module "$_type"; break
            done < /proc/mounts ||
            panic "could not copy root fs module"
        fi
    else
        print "copying all modules"

        find \
            "${moddir}/${kernel}/kernel/fs" \
            "${moddir}/${kernel}/kernel/lib" \
            "${moddir}/${kernel}/kernel/arch" \
            "${moddir}/${kernel}/kernel/crypto" \
            "${moddir}/${kernel}/kernel/drivers/md" \
            "${moddir}/${kernel}/kernel/drivers/ata" \
            "${moddir}/${kernel}/kernel/drivers/scsi" \
            "${moddir}/${kernel}/kernel/drivers/block" \
            "${moddir}/${kernel}/kernel/drivers/virtio" \
            "${moddir}/${kernel}/kernel/drivers/usb/host" \
            "${moddir}/${kernel}/kernel/drivers/usb/storage" \
            -type f 2> /dev/null |

        while read -r _module || [ "$_module" ]; do

            # strip path and extension
            _module="${_module##*/}"
            _module="${_module%%.*}"

            # skip unneeded modules and skip modules which
            # depends on them as well
            case $(modprobe -S "$kernel" -D "$_module") in
                *wmi* | *gpu* | *net*) continue ;;
            esac 2> /dev/null

            copy_module "$_module"
        done
    fi

    copy_binary modprobe

    cp "${moddir}/${kernel}/modules.builtin" \
       "${moddir}/${kernel}/modules.order" \
       "${tmpdir}${moddir}/${kernel}"

    chmod 644 \
        "${tmpdir}${moddir}/${kernel}/modules.builtin" \
        "${tmpdir}${moddir}/${kernel}/modules.order"

    depmod -b "$tmpdir" "$kernel"
}

copy_devmgr()
{
    print "configuring device manager"

    # false positive
    # shellcheck disable=2016
    case "$devmgr" in
        none)
            # TODO implement mode without device manager using deprecated
            # /sys/kernel/uevent_helper or /proc/sys/kernel/hotplug
        ;;
        mdev)
            for _binary in mdev find; do
                copy_binary "$_binary"
            done

            printf "%s\n" \
                'SUBSYSTEM=block;.* 0:0 660 @device-helper' \
            > "${tmpdir}/etc/mdev.conf"

            [ "$monolith" = 1 ] || printf "%s\n" \
                '$MODALIAS=.* 0:0 660 @modprobe "$MODALIAS"' \
            >> "${tmpdir}/etc/mdev.conf"
        ;;
        mdevd)
            for _binary in mdevd mdevd-coldplug; do
                copy_binary "$_binary"
            done

            printf "%s\n" \
                'SUBSYSTEM=block;.* 0:0 660 @device-helper' \
            > "${tmpdir}/etc/mdev.conf"

            [ "$monolith" = 1 ] || printf "%s\n" \
                '$MODALIAS=.* 0:0 660 @modprobe "$MODALIAS"' \
            >> "${tmpdir}/etc/mdev.conf"
        ;;
        udev)
            # why systemd violates FHS and places daemon in /lib ?
            udevd=$(command -v /lib/systemd/systemd-udevd) || udevd=udevd

            for _binary in "$udevd" udevadm; do
                copy_binary "$_binary"
            done

            mkdir -p "${tmpdir}/lib/udev/rules.d"

            printf "%s\n" \
                'SUBSYSTEMS=="block", ACTION=="add", RUN+="/bin/device-helper"' \
            > "${tmpdir}/lib/udev/rules.d/device-helper.rules"

            [ "$monolith" = 1 ] || printf "%s\n" \
                'ENV{MODALIAS}=="?*", ACTION=="add", RUN+="/bin/modprobe %E{MODALIAS}"' \
            >> "${tmpdir}/lib/udev/rules.d/device-helper.rules"
        ;;
    esac
}

make_initramfs()
(
    print "generating initramfs image"

    # check if image already exist
    [ "$force" != 1 ] && [ -e "$output" ] &&
        panic "initramfs image already exist"

    cd "$tmpdir"; find . |
    cpio -oH newc 2> /dev/null |
    ${compress:-cat} > "$output" ||
    panic "failed to generate initramfs image"

    print "done! check out $output"
)

# int main()
{
    [ "$(id -u)" = 0 ] || panic "must be run as root"

    # enable exit on error and disable globbing
    set -ef

    prepare_environment "$@"
    prepare_initramfs

    # copy and run hooks if any
    for _hook in $hooks; do
        copy_hook "$_hook"
    done

    copy_devmgr
    copy_modules
    make_initramfs
}