#!/bin/sh
#
# tiny initramfs generation tool

msg() {
    # print message
    case "$1" in
        info)
            printf "info >> %s\n" "$2" >&2
            ;;
        warn)
            printf "warning >> %s\n" "$2" >&2
            printf "do you want to continue? press enter or ctrl+c to exit\n"
            read -r _
            ;;
        panic)
            printf "panic >> %s\n" "$2" >&2
            exit 1
            ;;
    esac
}

usage() {
    cat << EOF
usage: $0 [options]
    -o, --output <file>     output file. default is $(readlink -f $(dirname "$0"))/initramfs-$(uname -r)
    -f, --files <dir>       files directory. default is $(readlink -f $(dirname "$0"))

EOF
}

parse_args() {
    while [ "$1" ]; do
        case "$1" in
            -o | --output)
                initramfs="$2"
                shift 2
                ;;
            -f | --files)
                filesdir="$2"
                shift 2
                ;;
            -h | --help)
                usage
                exit 0
                ;;
            *)
                printf "%s\n\n" "invalid option: '$1'"
                usage
                exit 1
                ;;
        esac
    done
}

create_wrkdir() {
    msg info "creating working directory"

    if [ "$XDG_CACHE_HOME" ]; then
        wrkdir="${XDG_CACHE_HOME}/initramfs.$$"
    elif [ "$TMPDIR" ]; then
        wrkdir="${TMPDIR}/initramfs.$$"
    else
        wrkdir="/tmp/initramfs.$$"
    fi

    mkdir "$wrkdir" || msg panic "failed to create working directory"
}

remove_wrkdir() {
    msg info "removing working directory"

    rm -rf "$wrkdir"
}

install_requirements() {
    msg info "installing requirements"

    # install user specified binaries
    [ "$binaries"       ] && install_binary "$binaries"

    # install util-linux binaries
    [ "$util_linux" = 1 ] && install_binary mount blkid

    # install mandatory binaries
    install_binary busybox modprobe
}

create_structure() {
    msg info "creating directory structure"

    for dir in dev tmp var run etc usr/lib usr/bin mnt/root proc root sys; do
        mkdir -p "${wrkdir}/${dir}"
    done
}

create_symlinks() {
    # some dynamically linked libraries and binaries compiled with hardcoded
    # dependencies path. to make it worked we need create symlinks for them.
    # also POSIX ln doesn't have --relative flag like in GNU ln. as workaround
    # we change directory to wrkdir and make needed symlinks.

    msg info "creating symlinks"

    ( cd "$wrkdir" && {
        ln -s usr/lib lib
        ln -s usr/lib lib64
        ln -s usr/bin bin
        ln -s usr/bin sbin
        ln -s ../run var/run
        cd "${wrkdir}/usr"
        ln -s bin sbin
        ln -s lib lib64
    } )
}

install_mdev() {
    msg info "installing mdev"

    install -m644  "${filesdir}/mdev.conf"      -t "${wrkdir}/etc"
    install -Dm755 "${filesdir}/storage-device" -t "${wrkdir}/lib/mdev"
}

install_mdevd() {
    msg info "installing mdevd"

    install_binary mdevd mdevd-coldplug

    install -m644  "${filesdir}/mdev.conf"      -t "${wrkdir}/etc"
    install -Dm755 "${filesdir}/storage-device" -t "${wrkdir}/lib/mdev"
}

install_udev() {
    msg info "installing udev"

    install_binary udevd udevadm dmsetup
    # FIXME rewrite this piece of crap
    find /usr/lib/udev -type f | grep -v "rc_keymaps\|hwdb.d" | cpio -pd "$wrkdir" > /dev/null 2>&1
}

install_lvm() {
    msg info "installing LVM"

    install_binary lvm

    lvm_drivers="dm-thin-pool dm-multipath dm-snapshot dm-cache dm-log dm-mirror"
    [ "$hostonly" = 1 ] && install_driver "$lvm_drivers"

    # install lvm config
    if [ "$lvm_conf" = 1 ]; then
        install -Dm644 /etc/lvm/*.conf -t "${wrkdir}/etc/lvm" || msg panic "failed to install LVM config"
    else
        mkdir "${wrkdir}/etc/lvm"
        cat << EOF > "${wrkdir}/etc/lvm/lvm.conf"
devices {
    # block discard support
    issue_discards = ${lvm_discard:-0}
}

global {
    # disable lvmetad
    use_lvmetad = 0
}
EOF
    fi
}

install_luks() {
    msg info "installing LUKS"

    install_binary cryptsetup

    luks_drivers="aes dm-crypt sha256 sha512 wp512 ecb lrw xts twofish serpent"
    [ "$hostonly" = 1 ] && install_driver "$luks_drivers"

    # avoid "locking directory missing" warning message
    mkdir "${wrkdir}/run/cryptsetup"

    # TODO get rid of this workaround
    # TODO fix static binary
    # workaround for luks2
    install -s -m755 /usr/lib/libgcc_s.so.1 -t "${wrkdir}/usr/lib" || msg panic "failed to install LUKS libraries"

    # copy luks header
    [ -f "$luks_header" ] && {
        install -m400 "$luks_header" "${wrkdir}/root/luks_header" || msg panic "failed to copy LUKS header"
        luks_args="--header=/root/luks_header $luks_args"
    }

    # copy luks keyfile
    [ -f "$luks_keyfile" ] && {
        install -m400 "$luks_keyfile" "${wrkdir}/root/luks_keyfile" || msg panic "failed to copy LUKS keyfile"
        luks_args="--key-file=/root/luks_keyfile $luks_args"
    }
}

install_driver() {
    # check monolithic kernel(builtin drivers)
    [ -d "${moddir}/${kernel}" ] || return

    # we need splitting
    # shellcheck disable=SC2068
    for driver in $@; do

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

        # TODO busybox modprobe doesn't support -S option
        modprobe -S "$kernel" -D "$driver" 2> /dev/null |
        grep -v "builtin\|net" |
        cut -d " " -f 2 |

        while read -r driver_dep; do
            install -Dm644 "$driver_dep" "${wrkdir}${driver_dep}"
        done
    done
}

install_hostonly_drivers() {
    msg info "installing hostonly drivers"

    [ "$root_type" ] || msg panic "hostonly mode required root_type option to be configured"

    # perform autodetection of drivers via /sys
    install_driver "$(find /sys -name modalias -exec sort -u "{}" "+")"

    # TODO autodetect root fs driver
    # TODO separate root type option
    # install root fs driver
    install_driver "$root_type"

    # install user specified drivers
    [ "$drivers" ] && install_driver "$drivers"
}

install_all_drivers() {
    msg info "installing all drivers"

    modker="${moddir}/${kernel}/kernel"

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

generate_depmod() {
    msg info "running depmod"

    modker="${moddir}/${kernel}"

    cp "${modker}/modules.builtin" "${modker}/modules.order" "${wrkdir}/${modker}"
    depmod -b "$wrkdir" "$kernel"
}

install_binary() {
    # TODO make strip optional

    # we need splitting
    # shellcheck disable=SC2068
    for binary in $@; do
        msg info "installing binary $binary"

        fullbin=$(command -v "$binary")

        # check binary existence
        [ "$fullbin" ] || msg panic "$binary doesn't exists"

        # install and strip binary
        install -s -m755 "$fullbin" -t "${wrkdir}/usr/bin"

        # check static
        ldd "$fullbin" > /dev/null 2>&1 || continue

        # install libraries
        install_library "$binary"
    done
}

install_library() {
    # TODO make strip optional

    wrkdirlib="${wrkdir}/usr/lib/"

    # extract paths to libraries from ldd output
    ldd $(command -v "$1") |
    sed -nr \
    -e 's,.* (/.*lib.*/.*.so.*) .*,\1,p' \
    -e 's,.*(/lib.*/ld.*.so.*) .*,\1,p' |

    while read -r library; do
        namelib="${library##*/}"
        reallib=$(readlink "$library")
        fulllib=$(readlink -f "$library")

        # check symlink
        if [ -h "$library" ]; then

            # check lib already existence
            [ -e "${wrkdirlib}${namelib}" ] ||
            [ -e "${wrkdir}${fulllib}"    ] ||
            {
                # regular
                install -s -m755 "${fulllib}" -t "${wrkdirlib}"

                # FIXME handle all symlinks
                # symlink may link to symlink
                [ -h "/usr/lib/${reallib}" ] &&
                cp -a "/usr/lib/${reallib}" "${wrkdirlib}"

                # symlink
                cp -a "$library" "${wrkdirlib}"
            }
        else
            [ -e "${wrkdirlib}${namelib}" ] ||
            install -s -m755 "$library" -t "${wrkdirlib}"
        fi
    done
}

install_files() {
    msg info "installing files"

    cat << EOF > "${wrkdir}/config"
debug="$debug"
init="$init"
root="$root"
root_type="$root_type"
root_opts="$root_opts"
devmgr="$devmgr"
#drivers
lvm="$lvm"
lvm_name="$lvm_name"
lvm_group="$lvm_group"
#lvm_discard
lvm_args="$lvm_args"
luks="$luks"
luks_root="$luks_root"
luks_name="$luks_name"
#luks_header
#luks_keyfile
luks_discard="$luks_discard"
luks_args="$luks_args"
EOF

    # needed for devmgr
    cat << EOF > "${wrkdir}/etc/group"
root:x:0:
tty:x:5:
dialout:x:11:
uucp:x:14:
kmem:x:3:
input:x:25:
video:x:13:
audio:x:12:
lp:x:10:
disk:x:9:
cdrom:x:16:
tape:x:6:
kvm:x:24:
floppy:x:8:
EOF

    # needed for devmgr
    cat << EOF > "${wrkdir}/etc/passwd"
root:x:0:0::/root:/bin/sh
nobody:x:99:99::/:/bin/false
EOF

    # install init script
    install -m755 "${filesdir}/init" -t "$wrkdir"
}

create_initramfs() {
    msg info "creating initramfs image"
    {
        (
            cd "$wrkdir"
            find . | cpio -oH newc | ${compress:-gzip -9}

        ) | tee "${initramfs:-${filesdir}/initramfs-${kernel}}"

    } > /dev/null 2>&1 || msg panic "failed to generate initramfs image"
}

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

parse_args "$@"

. "${filesdir:=$(readlink -f $(dirname "$0"))}/config" || msg panic "failed to source config"

# remove wrkdir on exit or unexpected error
trap remove_wrkdir EXIT INT

[ "$debug" = 1 ] && {
    # debug shell commands
    set -x
    # don't remove anything
    trap - EXIT INT
}

kernel="${kernel:-$(uname -r)}"
moddir="/lib/modules"

create_wrkdir
create_structure
create_symlinks
install_requirements

if [ "$hostonly" = 1 ]; then
    install_hostonly_drivers
else
    install_all_drivers
fi

generate_depmod

case "$devmgr" in
    udev)  install_udev ;;
    mdev)  install_mdev ;;
    mdevd) install_mdevd ;;
    *)     msg panic "devmgr option broken" ;;
esac

[ "$lvm"  = 1 ] && [ -x "$(command -v lvm)"        ] && install_lvm
[ "$luks" = 1 ] && [ -x "$(command -v cryptsetup)" ] && install_luks
install_files
create_initramfs

msg info "done! check out ${initramfs:-${filesdir}/initramfs-${kernel}}"