#!/bin/sh -ef
#
# tiny initramfs
#
# false positive
# shellcheck disable=2154

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

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

usage()
{
    cat << EOF
usage: $0 [option]
       -o, --output <file>  set initramfs output path
       -c, --config <file>  set config file path
       -m, --moddir <dir>   set modules directory
       -k, --kernel <ver>   set kernel version
       -F, --files  <dir>   set files directory
       -d, --debug          enable debug mode
       -f, --force          overwrite initramfs image

EOF
}

parse_args()
{
    while [ "$1" ]; do case "$1" in
        -o | --output)
            _output="${2:?}"
            shift 2
        ;;
        -c | --config)
            _config="${2:?}"
            shift 2
        ;;
        -m | --moddir)
            _moddir="${2:?}"
            shift 2
        ;;
        -k | --kernel)
            _kernel="${2:?}"
            shift 2
        ;;
        -F | --files)
            _filesdir="${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
}

prepare_environment()
{
    print "preparing environment"

    # false positive
    # shellcheck disable=1090
    for _file in $_config /etc/tinyramfs/config ./config; do
        [ -f "$_file" ] && { . "$_file"; break; }
    done || panic "failed to source config"

    for _dir in $_filesdir /usr/share/tinyramfs .; do
        [ -d "$_dir" ] && { filesdir="$_dir"; break; }
    done || panic "failed to locate required files"

    # general variables
    debug="${_debug:-${debug:-0}}"
    force="${_force:-${force:-0}}"
    moddir="${_moddir:-${moddir:-/lib/modules}}"
    kernel="${_kernel:-${kernel:-$(uname -r)}}"
    output="${_output:-${output:-/tmp/initramfs-${kernel}}}"

    mkdir -p "${workdir=${XDG_CACHE_HOME:-${TMPDIR:-/tmp}}/initramfs.$$}" ||
        panic "failed to create working directory"

    # helpers variables
    workdirbin="${workdir}/usr/bin/"
    workdirlib="${workdir}/usr/lib/"
    modker="${moddir}/${kernel}"
    OLD_IFS="$IFS"

    trap trap_helper EXIT INT

    # false positive
    # shellcheck disable=2015
    [ "$debug" = 1 ] && set -x ||:
}

trap_helper()
{
    # TODO need cleanup
    ret="$?"

    trap - EXIT INT

    [ "$debug" != 1 ] && {
        print "removing working directory"
        rm -rf "$workdir"
    }

    [ "$ret" != 0 ] && panic "something went wrong"
}

install_requirements()
{
    print "installing requirements"

    # install user specified and required binaries
    for _binary in \[ sh ln sleep mount printf setsid switch_root $binaries; do
        install_binary "$_binary"
    done

    # copy init
    install -m755 "${filesdir}/init" "${workdir}/init"

    # copy config
    printf "%s\n" \
        root="$root" \
        root_type="$root_type" \
        root_opts="$root_opts" \
        devmgr="$devmgr" \
        monolith="$monolith" \
    >> "${workdir}/etc/config"
}

create_structure()
{
    print "creating directory structure"

    mkdir -p \
        "${workdir}/etc" \
        "${workdir}/dev" \
        "${workdir}/sys" \
        "${workdir}/tmp" \
        "${workdir}/proc" \
        "${workdir}/root" \
        "${workdir}/usr/lib" \
        "${workdir}/usr/bin" \
        "${workdir}/mnt/root"
}

create_symlinks()
{
    print "creating symlinks"

    ln -s usr/lib "${workdir}/lib"
    ln -s usr/lib "${workdir}/lib64"
    ln -s usr/bin "${workdir}/bin"
    ln -s usr/bin "${workdir}/sbin"
    ln -s bin     "${workdir}/usr/sbin"
    ln -s lib     "${workdir}/usr/lib64"
}

install_devmgr()
{
    print "installing device manager"

    install_device_helper()
    {
        for _binary in kill mkdir blkid "${filesdir}/device-helper"; do
            install_binary "$_binary"
        done

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

        # false positive
        # shellcheck disable=2016
        [ "$monolith" != 1 ] && printf "%s\n" \
            '$MODALIAS=.*       0:0 660 @modprobe "$MODALIAS"' \
        >> "${workdir}/etc/mdev.conf"
    }

    # TODO investigate booting without device manager
    case "$devmgr" in
        udev)
            for _binary in udevd udevadm; do
                install_binary "$_binary"
            done

            # i hate udev
            mkdir -p "${workdir}/run/udev"

            # exclusively handle requirement
            [ "$luks" = 1 ] || [ "$lvm" = 1 ] && install_binary dmsetup

            for _binary in /usr/lib/udev/ata_id /usr/lib/udev/scsi_id; do
                install -Dm755 "$_binary" "${workdir}${_binary}" || return
            done

            # TODO we really need all rules?
            set +f; for _file in /usr/lib/udev/rules.d/*; do
                install -Dm644 "$_file" "${workdir}${_file}" || return
            done; set -f
        ;;
        mdev)
            install_binary mdev
            install_device_helper
        ;;
        mdevd)
            for _binary in mdevd mdevd-coldplug; do
                install_binary "$_binary"
            done

            install_device_helper
        ;;
    esac
}

install_lvm()
{
    print "installing LVM"

    for _binary in lvchange vgchange; do
        install_binary "$_binary"
    done

    lvm_config="
        devices {
            write_cache_state = 0
        }
        backup {
            backup  = 0
            archive = 0
        }
        global {
            use_lvmetad = 0
        }"

    # word splitting is safe by design
    # shellcheck disable=2086
    { IFS=,; set -- $lvm_opts; IFS="$OLD_IFS"; }

    for opt; do case "$opt" in
        config | config=1)
            embed_lvm_config=1
        ;;
    esac; done

    mkdir -p "${workdir}/etc/lvm"; lvmconfig \
        --config "$lvm_config" \
        ${embed_lvm_config:+--mergedconfig} \
    > "${workdir}/etc/lvm/lvm.conf"

    # copy config
    printf "%s\n" \
        lvm="$lvm" \
        lvm_opts="$lvm_opts" \
    >> "${workdir}/etc/config"
}

install_luks()
{
    print "installing LUKS"

    install_binary cryptsetup

    # avoid libgcc_s.so.1 missing error
    # see https://bugs.archlinux.org/task/56771
    [ -e /usr/lib/libgcc_s.so.1 ] && install_library /usr/lib/libgcc_s.so.1

    # word splitting is safe by design
    # shellcheck disable=2086
    { IFS=,; set -- $luks_opts; IFS="$OLD_IFS"; }

    for opt; do case "${opt%%=*}" in
        header)
            install -m400 "${opt##*=}" "${workdir}/root/header" || return
            luks_opts=$(printf "%s" "$luks_opts" | sed "s|${opt##*=}|/root/header|")
        ;;
        key)
            install -m400 "${opt##*=}" "${workdir}/root/key" || return
            luks_opts=$(printf "%s" "$luks_opts" | sed "s|${opt##*=}|/root/key|")
        ;;
    esac; done

    # copy config
    printf "%s\n" \
        luks="$luks" \
        luks_root="$luks_root" \
        luks_opts="$luks_opts" \
    >> "${workdir}/etc/config"
}

install_module()
{
    module="$1"

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

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

        # strip unneeded stuff
        for _exclude_module in wmi gpu net sound builtin $modules_exclude; do
            case "$module" in *"$_exclude_module"*) continue 2 ;; esac
        done

        module="${module#insmod }"

        # check if module already installed
        [ -e "$module" ] && [ ! -e "${workdir}${module}" ] &&
            install -Dm644 "$module" "${workdir}${module}" || return
    done ||:
}

install_hostonly_modules()
{
    print "installing hostonly modules"

    # perform autodetection of modules via /sys
    find /sys -name modalias -exec sort -u {} + |

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

    # install LVM modules
    [ "$lvm" = 1 ] &&
        for _module in dm-thin-pool dm-multipath dm-snapshot dm-cache dm-log dm-mirror; do
            install_module "$_module"
        done

    # install LUKS modules
    [ "$luks" = 1 ] &&
        for _module in aes dm-crypt sha256 sha512 wp512 ecb lrw xts twofish serpent; do
            install_module "$_module"
        done

    # install root partition module
    if [ "$root_type" ]; then
        install_module "$root_type"
    else
        while read -r _ _dir _type _ _ _ || [ "$_dir" ]; do
            [ "$_dir" = / ] && {
                install_module "${root_type=$_type}"
                break
            }
        done < /proc/mounts || panic "failed to install root partition module"
    fi
}

install_all_modules()
{
    print "installing all modules"

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

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

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

        install_module "$_module"
    done ||:
}

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

    # check if binary exist and builtin
    case "$binary" in
        */*)
            : no operation
        ;;
        "")
            panic "$1 doesn't exist"
        ;;
        *)
            # word splitting is safe by design
            # shellcheck disable=2086
            { IFS=:; set -- $PATH; IFS="$OLD_IFS"; }

            for _dir; do
                [ -x "${_dir}/${binary}" ] && {
                    binary="${_dir}/${binary}"
                    break
                }
            done || panic "couldn't find external $1 binary"
        ;;
    esac

    # check if binary already installed
    [ -e "${workdirbin}${binary##*/}" ] && return

    # iterate throught symlinks and copy them
    while [ -h "$binary" ]; do
        cp -P "$binary" "$workdirbin"
        readlink_binary=$(readlink "$binary")
        binary="${binary%/*}/${readlink_binary##*/}"
    done

    install -m755 "$binary" "${workdirbin}${binary##*/}" || return
    strip "${workdirbin}${binary##*/}" > /dev/null 2>&1 ||:

    # check if binary statically linked
    ldd "$binary" > /dev/null 2>&1 || return 0

    # exract paths to libraries
    ldd "$binary" |

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

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

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

        [ -e "$_library" ] && install_library "$_library"
    done ||:
}

install_library()
{
    library="$1"

    # check if library already installed
    [ -e "${workdirlib}${library##*/}" ] && return

    # iterate throught symlinks and copy them
    while [ -h "$library" ]; do
        cp -P "$library" "$workdirlib"
        readlink_library=$(readlink "$library")
        library="${library%/*}/${readlink_library##*/}"
    done

    install -m755 "$library" "${workdirlib}${library##*/}" || return
    strip "${workdirlib}${library##*/}" > /dev/null 2>&1 ||:
}

create_initramfs()
{
    print "creating initramfs image"

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

    (
        cd "$workdir"; find . |

        if [ "$compress" = none ]; then
            cpio -oH newc
        else
            cpio -oH newc | ${compress:-gzip -9}
        fi

    ) > "$output" 2> /dev/null ||
        panic "failed to generate initramfs image"
}

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

    parse_args "$@"
    prepare_environment
    create_structure
    create_symlinks

    [ "$lvm"  = 1 ] && install_lvm
    [ "$luks" = 1 ] && install_luks

    # check monolithic kernel
    [ "$monolith" != 1 ] && [ -d "$moddir" ] && {

        # check hostonly mode
        if [ "$hostonly" = 1 ]; then
            install_hostonly_modules
        else
            install_all_modules
        fi

        # install user specified modules if any
        [ "$modules" ] && {
            for _module in $modules; do
                install_module "$_module"
            done

            # copy config
            printf "%s\n" \
                modules="$modules" \
            >> "${workdir}/etc/config"
        }

        for _binary in find sort modprobe; do
            install_binary "$_binary"
        done

        install -m644 \
            "${modker}/modules.builtin" \
            "${modker}/modules.order" \
            "${workdir}${modker}"

        depmod -b "$workdir" "$kernel"
    }

    install_devmgr
    install_requirements
    create_initramfs

    print "done! check out $output"
}