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

# print message
msg() {
    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
}

# create tmpdir
create_tmpdir() {
    if [ -n "$XDG_CACHE_HOME" ]; then
        tmpdir="${XDG_CACHE_HOME}/initramfs.$$"
        mkdir "$tmpdir"
        printf "%s\n" "$tmpdir"
    elif [ -n "$TMPDIR" ]; then
        tmpdir="${TMPDIR}/initramfs.$$"
        mkdir "$tmpdir"
        printf "%s\n" "$tmpdir"
    else
        tmpdir="/tmp/initramfs.$$"
        mkdir "$tmpdir"
        printf "%s\n" "$tmpdir"
    fi
}

# remove tmpdir
remove_tmpdir() {
    msg info "removing work dir"
    rm -rf "$tmpdir"
}

# change current directory to script directory if user haven't do it
check_currentdir() {
    script_dir=$(dirname $(readlink -f "$0"))
    [ "$PWD" = "$script_dir" ] || {
        msg info "changing directory to script dir"
        cd "$script_dir" || msg panic "failed to change directory"
    }
}

# check needed files
check_requirements() {
    msg info "checking requirements"

    # check busybox installed
    command -v busybox > /dev/null 2>&1 && {
        # check busybox supports CONFIG_FEATURE_INSTALLER
        busybox --help | grep -q "\-\-install \[\-s\]" || msg panic "recompile busybox with CONFIG_FEATURE_INSTALLER"

        # check requirements for init
        for busybox_dep in mdev uevent switch_root; do
            busybox $busybox_dep --help > /dev/null 2>&1 || msg panic "recompile busybox with $busybox_dep"
        done
    } || msg panic "busybox doesn't installed"

    # check kmod modprobe installed
    command -v modprobe > /dev/null 2>&1 && {
        # busybox modprobe doesn't supported(yet)
        modprobe --version 2>&1 | grep -q "kmod" || msg panic "kmod modprobe version doesn't installed"
    } || msg panic "modprobe doesn't installed"

    # check util-linux tools
    [ "$use_util_linux" = 1 ] && {
        # check mount installed
        if command -v mount > /dev/null 2>&1; then
            mount --version 2>&1 | grep -q "util-linux" || {
                msg warning "util-linux mount version doesn't installed. PARTUUID and filesystem type autodetection support will be missing"
                use_util_linux=0
            }
        else
            msg panic "mount doesn't installed"
        fi

        # check blkid installed
        if command -v blkid > /dev/null 2>&1; then
            blkid --version 2>&1 | grep -q "util-linux" || {
                msg warning "util-linux blkid version doesn't installed. PARTUUID support will be missing"
                use_util_linux=0
            }
        else
            msg panic "blkid doesn't installed"
        fi
    }
}

# install requirements
install_requirements() {
    msg info "installing requirements"

    # install user specified binaries
    [ -n "$binaries" ] && install_binaries $binaries

    # install util-linux binaries
    [ "$use_util_linux" = 1 ] && install_binaries mount blkid

    # install mandatory binaries
    install_binaries busybox modprobe
}

# create FHS directory structure
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 "${tmpdir}/${dir}"
    done
}

# 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 tmpdir and make needed symlinks.
create_symlinks() {
    msg info "creating symlinks"
    ( cd "$tmpdir" && {
        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 "${tmpdir}/usr"
        ln -s bin sbin
        ln -s lib lib64
    } )
}

#parse_fstab() {
# TODO parse fstab
#while [ "$use_fstab" -eq 1 ] && read fs dir type opts; do thing; done < /etc/fstab
#}

#parse_crypttab() {
# TODO parse crypttab
#}

# install mdev
install_mdev() {
    msg info "installing mdev"
    install -m644 mdev.conf -t "${tmpdir}/etc"
    install -Dm755 storage-device -t "${tmpdir}/lib/mdev"
}

# install mdevd
install_mdevd() {
    msg info "installing mdevd"
    install_binaries mdevd mdevd-coldplug
    install -m644 mdev.conf -t "${tmpdir}/etc"
    install -Dm755 storage-device -t "${tmpdir}/lib/mdev"
}

# install udev
install_udev() {
    msg info "installing udev"
    install_binaries udevd udevadm dmsetup
    # FIXME rewrite this piece of crap
    find /usr/lib/udev -type f | grep -v "rc_keymaps\|hwdb.d" | cpio -pd "$tmpdir" > /dev/null 2>&1
}

# handle lvm
install_lvm() {
    msg info "installing LVM"
    install_binaries lvm

    # if hostonly mode enabled install only needed drivers
    [ "$hostonly" = 1 ] && {
        for lvm_driver in dm-thin-pool dm-multipath dm-snapshot dm-cache dm-log dm-mirror; do
            for lvm_driver_dep in $(modprobe -D "$lvm_driver" 2> /dev/null | grep -v builtin | cut -d " " -f 2); do
                install -Dm644 "$lvm_driver_dep" "${tmpdir}${lvm_driver_dep}"
            done
        done
    }

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

global {
    use_lvmetad = 0
}
EOF
    fi
}

# handle luks
install_luks() {
    msg info "installing LUKS"
    install_binaries cryptsetup

    # if hostonly mode enabled install only needed drivers
    [ "$hostonly" = 1 ] && {
        for luks_driver in aes dm-crypt sha256 sha512 wp512 ecb lrw xts twofish serpent; do
            for luks_driver_dep in $(modprobe -D "$luks_driver" 2> /dev/null | grep -v builtin | cut -d " " -f 2); do
                install -Dm644 "$luks_driver_dep" "${tmpdir}${luks_driver_dep}"
            done
        done
    }

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

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

    # block discard support
    [ "$luks_discard" = 1 ] && luks_args="--allow-discards $luks_args"

    # copy luks header
    [ -f "$luks_header" ] && {
        install -m400 "$luks_header" "${tmpdir}/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" "${tmpdir}/root/luks_keyfile" || msg panic "failed to copy LUKS keyfile"
        luks_args="--key-file=/root/luks_keyfile $luks_args"
    }
}

# install drivers
install_drivers() {
    msg info "installing hostonly drivers"
    [ -n "$root_type" ] || msg panic "hostonly mode required root_type option to be configured"

    # perform autodetection of drivers via /sys
    find /sys/devices -name modalias -exec sort -u "{}" "+" | while read -r driver; do
        for driver_dep in $(modprobe -D "$driver" 2> /dev/null | grep -v builtin | cut -d " " -f 2); do
            install -Dm644 "$driver_dep" "${tmpdir}${driver_dep}"
        done
    done

    # TODO autodetect root fs driver
    # TODO separate root type option
    # install root fs driver
    for root_driver in $(modprobe -D "$root_type" 2> /dev/null | grep -v builtin | cut -d " " -f 2); do
        install -Dm644 "$root_driver" "${tmpdir}${root_driver}"
    done

    # install user specified drivers
    [ -n "$drivers" ] && {
        printf "%s\n" "$drivers" | while read -r custom_driver; do
            for custom_driver_dep in $(modprobe -D "$custom_driver" 2> /dev/null | grep -v builtin | cut -d " " -f 2); do
                install -Dm644 "$custom_driver_dep" "${tmpdir}${custom_driver_dep}"
            done
        done
    }
}

# find and install all drivers
install_all_drivers() {
    msg info "installing all drivers"
    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 | cpio -pd "$tmpdir" > /dev/null 2>&1
}

# generate "modules" files
generate_depmod() {
    msg info "generating dependendies list of drivers"
    # install list of drivers
    cp "${modker}/modules.softdep" "${modker}/modules.builtin" "${modker}/modules.order" "${tmpdir}/${modker}"

    # generate dependencies list of drivers
    depmod -b "$tmpdir" "$kernel"
}

# TODO make strip optional
# handle binaries
install_binaries() {
    printf "%s\n" "$@" | while read -r binary; do
        msg info "installing binary $binary"
        # check binary existence
        command -v "$binary" > /dev/null 2>&1 || msg panic "$binary doesn't exists"

        # install and strip binary
        install -s -m755 "$(command -v $binary)" -t "${tmpdir}/usr/bin"

        # check statically linking
        ldd "$(command -v $binary)" > /dev/null 2>&1 || continue

        # install libraries
        install_libraries $binary
    done
}

# TODO make strip optional
# handle libraries
install_libraries() {
    for library in $(ldd "$(command -v $1)" | sed -nre 's,.* (/.*lib.*/.*.so.*) .*,\1,p' -e 's,.*(/lib.*/ld.*.so.*) .*,\1,p'); do
        msg info "installing library $library"
        # check symlink
        if [ -h "$library" ]; then
            # check lib already existence
            if [ ! -e "${tmpdir}/usr/lib/${library##*/}" ] && [ ! -e "${tmpdir}/$(readlink -f $library)" ]; then
                # regular
                install -s -m755 "$(readlink -f $library)" -t "${tmpdir}/usr/lib"

                # FIXME handle all symlinks
                # symlink may link to symlink
                [ -h "/usr/lib/$(readlink $library)" ] && cp -a "/usr/lib/$(readlink $library)" "${tmpdir}/usr/lib"

                # symlink
                cp -a "$library" "${tmpdir}/usr/lib"
            fi
        else
            if [ ! -e "${tmpdir}/usr/lib/${library##*/}" ]; then
                install -s -m755 "$library" -t "${tmpdir}/usr/lib"
            fi
        fi
    done
}

# install important files used by init
install_files() {
    msg info "installing files"
    # FIXME eof broken
    # initialize config
    cat << EOF > "${tmpdir}/config"
debug="$debug"
root="$root"
root_type="$root_type"
root_args="$root_args"
devmgr="$devmgr"
#drivers="$drivers"
use_lvm="$use_lvm"
use_luks="$use_luks"
luks_root="$luks_root"
luks_args="$luks_args"
EOF

    # needed for devmgr
    cat << EOF > "${tmpdir}/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 > "${tmpdir}/etc/passwd"
root:x:0:0::/root:/bin/sh
nobody:x:99:99::/:/bin/false
EOF

    # install init script
    install -m755 ./init -t "$tmpdir"
}

# TODO add more compession tools
# create and compress cpio archive
create_initramfs() {
    msg info "creating initramfs image"
    {
        (
            cd "$tmpdir"
            find . \
                | cpio -oH newc \
                | gzip -9
        ) \
            | tee "${script_dir}/initramfs-${kernel}.img.gz"
    } > /dev/null 2>&1 || msg panic "failed to generate initramfs image"
}

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

# remove tmpdir on exit or unexpected error
trap remove_tmpdir EXIT INT

check_currentdir

# source config
. ./config || msg panic "./config doesn't exists"

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

# variables
tmpdir="$(create_tmpdir)" || msg panic "failed to create working directory"
kernel="$(uname -r)"
moddir="/lib/modules"
modker="${moddir}/${kernel}"

check_requirements
create_structure
create_symlinks
#parse_fstab
#parse_crypttab
install_requirements

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

generate_depmod

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

[ "$use_luks" = 1 ] && install_luks
[ "$use_lvm" = 1 ] && install_lvm

install_files
create_initramfs

msg info "done! check out initramfs-${kernel}.img.gz"