#!/bin/bash

# Copyright (C) 2012-2020 Jonathan Vasquez <jon@xyinn.org>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ========== Variables ==========
_use_encryption=0
_zfs_encryption_key=""
_new_root="/mnt/root"
_init="/sbin/init"
_cache_file="/etc/zfs/zpool.cache"
_rootfs_cache_file="${_new_root}${_cache_file}"
_bliss_version_file="/version.bliss"
_bliss_modules_file="/modules.bliss"

# Starts a rescue shell
RescueShell()
{
    Warn "Booting into rescue shell..." && NewLine
    setsid sh -c 'exec sh < /dev/tty1 > /dev/tty1 2>&1'
}

# Loads Bliss Version (Information Purposes Only)
LoadVersion()
{
    _version=$(cat ${_bliss_version_file})
}

# Module loading function
LoadModules()
{
    Info "Loading modules..."

    IFS="," read -a _modules <<< $(cat ${_bliss_modules_file})

    for module in ${_modules[@]}; do
        modprobe ${module}
    done
}

# Cleanly mounts the required devices
MountRequiredDevices()
{
    Info "Mounting kernel devices..."

    mount -t proc none /proc
    mount -t devtmpfs none /dev
    mount -t sysfs none /sys
    mount -t tmpfs -o mode=755,nodev none /run
}

# Cleanly umounts the required devices
UnmountRequiredDevices()
{
    Info "Unmounting kernel devices..."

    umount -l /proc
    umount -l /dev
    umount -l /sys
}

# Mounts /run onto the rootfs before switching
PremountRunOnNewRoot()
{
    Info "Premounting /run onto the rootfs..."

    # Mounts /run onto the rootfs before switching. Mounting this is a requirement
    # to implement the systemd initrd interface. Without this, I've experienced
    # that the tasks (i.e luks decrypted drives) would actually stall systemd
    # for 1m30s if they were listed in /etc/fstab.

    # Link: http://www.freedesktop.org/wiki/Software/systemd/InitrdInterface/
    local targetRunDirectory="${_new_root}"/run
    mount --rbind /run "${targetRunDirectory}" || Fail "Failed to bind /run into the rootfs!"
}

# Loads the user's keymap if it exists
LoadKeymapIfNeeded()
{
    local path_to_keymap="/etc/keymap"

    if [[ -f ${path_to_keymap} ]]; then
        Info "Loading keymap..."
        loadkeys "${path_to_keymap}" > /dev/null 2>&1
    fi
}

# Parses an individual parameter from the command line
ParseOption()
{
    echo "${1#*=}"
}

# Starts udev and udevadm
StartUdev()
{
    Info "Starting udev..."
    udevd --daemon --resolve-names=never 2> /dev/null
    GenerateDeviceLinks
}

# Starts udevadm and generates the device symlinks (uuid, label, etc)
GenerateDeviceLinks()
{
    udevadm trigger
    udevadm settle
}

# Stops udev from running so that we don't have problems when systemd attempts to run udev itself
StopUdev()
{
    Info "Stopping udev..."
    killall udevd
}

# Process command line parameters
ParseKernelParameters()
{
    Info "Parsing kernel parameters..."

    for param in $(</proc/cmdline); do
        case "${param}" in
            root=*)
                _root=$(RetrieveDriveValue "${param}")
                ;;
            options=*)
                _options=$(ParseOption "${param}")
                ;;
            usr=*)
                _usr=$(ParseOption "${param}")
                ;;
            encrypted)
                _use_encryption=1
                ;;
            init=*)
                _init=$(ParseOption "${param}")
                ;;
            by=*)
                _by=$(ParseByOption "${param}")
                ;;
            refresh)
                _refresh=1
                ;;
            recover)
                _recover=1
                ;;
            su)
                _su=1
                ;;
        esac
    done

    if [[ -z ${_root} ]]; then
        Fail "You must pass the 'root' variable."
    fi
}

# Parses and sets the 'by' value so that we can have
# a more sanitized path for scanning the devices directories
# later on when we try to mount our zpool.
ParseByOption()
{
    if [[ -z $1 ]]; then
        Fail "No 'by' value was passed into this function!"
    fi

    # Set the following shell property so that we can
    # match regardless of case. This property will be unset
    # at the end of the function.
    shopt -s nocasematch

    local option=$(ParseOption "$1")

    case "${option}" in
        dev)
            echo "/dev"
            ;;
        id)
            echo "/dev/disk/by-id"
            ;;
        uuid)
            echo "/dev/disk/by-uuid"
            ;;
        partuuid)
            echo "/dev/disk/by-partuuid"
            ;;
        label)
            echo "/dev/disk/by-label"
            ;;
        partlabel)
            echo "/dev/disk/by-partlabel"
            ;;
        *)
            echo "${option}"
            ;;
    esac

    shopt -u nocasematch
}

# Retrieves the proper value for the drive
RetrieveDriveValue()
{
    if [[ -z $1 ]]; then
        Fail "No drive was passed into the function!"
    fi

    local tempDrive=$(ParseOption "$1")

    case "${tempDrive}" in
        ID=*)
            eval "${tempDrive}"
            echo "/dev/disk/by-id/${ID}"
            ;;
        UUID=*)
            eval "${tempDrive}"
            echo "/dev/disk/by-uuid/${UUID}"
            ;;
        PARTUUID=*)
            eval "${tempDrive}"
            echo "/dev/disk/by-partuuid/${PARTUUID}"
            ;;
        LABEL=*)
            eval "${tempDrive}"
            echo "/dev/disk/by-label/${LABEL}"
            ;;
        PARTLABEL=*)
            eval "${tempDrive}"
            echo "/dev/disk/by-partlabel/${PARTLABEL}"
            ;;
        *)
            echo "${tempDrive}"
            ;;
    esac
}

# Gets a decryption key without displaying it on screen. Retrieving the
# decryption key will test and block until a correct key is given.
GetDecryptionKey()
{
    _was_pool_imported=0

    while [[ -z "${_zfs_encryption_key}" ]]; do
        Ask "Enter passphrase: " && read -s -r _zfs_encryption_key && NewLine

        if [[ ${_was_pool_imported} -eq 0 ]]; then
            PrepForSafeZfsCommand && zpool import -f -o readonly=on "${_pool_name}" -R /tmp/random_dir
            _was_pool_imported=1
        fi

        LoadZfsKey "dry-run"

        [[ $? -ne 0 ]] && unset _zfs_encryption_key
    done

    [[ ${_was_pool_imported} -eq 1 ]] && PrepForSafeZfsCommand && zpool export "${_pool_name}"
}

# Loads the encryption key for the pool
LoadZfsKey()
{
    if [[ $1 == "dry-run" ]]; then
        echo "${_zfs_encryption_key}" | zfs load-key -n "${_pool_name}"
    else
        echo "${_zfs_encryption_key}" | zfs load-key "${_pool_name}"
    fi
}

# Retrieves the zpool.cache file from the rootfs
GetZpoolCacheFromSystem()
{
    Info "Retrieving zpool.cache from \"${_pool_name}\"..."

    if [[ -z "${_by}" ]]; then
        PrepForSafeZfsCommand && zpool import -f -N -o readonly=on "${_pool_name}"
    else
        PrepForSafeZfsCommand && zpool import -d "${_by}" -f -N -o readonly=on "${_pool_name}"
    fi

    [[ ${_use_encryption} -eq 1 ]] && LoadZfsKey

    PrepForSafeZfsCommand && mount -t zfs -o ro,zfsutil "${_root}" "${_new_root}"

    if [[ -f "${_rootfs_cache_file}" ]]; then
        cp "${_rootfs_cache_file}" /etc/zfs
        if [[ ! -f "${_cache_file}" ]]; then
            Warn "Unable to copy zpool.cache from \"${_pool_name}\"!"
        fi
    else
        Warn "zpool.cache file does not exist on \"${_pool_name}\"."
        Warn "It will automatically be copied into your system."
    fi

    # Let's also call the safe command before a normal unmount
    # (since the unmount is ultimately acting on a zfs dataset)
    PrepForSafeZfsCommand && umount "${_new_root}"

    # Use -F (undocumented flag) so that the copied zpool.cache is not nuked.
    PrepForSafeZfsCommand && zpool export -F "${_pool_name}"
}

# Run ZFS Specific Code
ZfsTrigger()
{
    _pool_name="${_root%%/*}"

    Flag "Importing \"${_pool_name}\"..."

    # Request the encryption key if needed. This will be cached and re-used
    # as we attempt to boot into the system. For now we are only supporting
    # "passphrase based decryption". The passphrase needed will be the
    # passphrase set at the pool level, so make sure that your OS dataset is
    # using that same one if multiple passphrases are being used.
    if [[ ${_use_encryption} -eq 1 ]]; then
        Info "Loading encryption key for \"${_pool_name}\"..."
        GetDecryptionKey
    fi

    # Use the system's cache file only if we don't want to reset it.
    if [[ ${_refresh} -ne 1 ]]; then
        GetZpoolCacheFromSystem
    fi

    if [[ -f "${_cache_file}" ]]; then
        # In order for this to be successful, the paths for each device in the
        # zpool's cache file need to be in the initramfs exactly as they were
        # when the pool was last explictly imported. In other words, the paths
        # in the zpool.cache for each drive must exist in the initramfs.
        Info "Loading \"${_pool_name}\" from zpool.cache..."

        # We are going to use the zpool.cache file but we will only mount from
        # the pool that contains our root dataset. Post early boot should be
        # responsible for mounting the rest of the pools/datasets. If we
        # attempt to mount all the pools at this point, there is a possibility
        # that the boot will fail (Example: You had an encrypted backup
        # drive/pool enabled, then you rebooted your machine, and the initramfs
        # fails because the backup pool is listed in the cache, but all the
        # drives are locked.)
        PrepForSafeZfsCommand && zpool import -N -c "${_cache_file}" "${_pool_name}"

        if [[ $? -eq 0 ]]; then
            [[ ${_use_encryption} -eq 1 ]] && LoadZfsKey
            return
        fi

        Warn "Unable to load your cache file. Falling back to auto-detection."

        # Remove the existing cache file since it's useless at this point
        rm "${_cache_file}"
    fi

    # If the user specified the "by-<>" type that they would like to
    # use, then attempt to mount the pool from that directory.
    if [[ ! -z ${_by} ]]; then
        PrepForSafeZfsCommand && zpool import -d "${_by}" -f -N -o cachefile= "${_pool_name}" || \
            Fail "Failed to mount \"${_pool_name}\" from the \"${_by}\" directory."

        [[ ${_use_encryption} -eq 1 ]] && LoadZfsKey

        Info "Successfully mounted \"${_pool_name}\" using the \"${_by}\" directory."
    else
        _device_dirs=(
            /dev
            /dev/mapper
            /dev/disk/by-*
        )

        local success=1

        for i in ${!_device_dirs[@]}; do
            local dir=${_device_dirs[i]}

            PrepForSafeZfsCommand && zpool import -d "${dir}" -f -N -o cachefile= "${_pool_name}"

            if [[ $? == 0 ]]; then
                Info "Successfully mounted \"${_pool_name}\" using the \"${dir}\" directory."
                success=0
                break
            fi
        done

        if [[ ${success} -ne 0 ]]; then
            Fail "Failed to mount \"${_pool_name}\" after scanning all the device directories."
        fi

        [[ ${_use_encryption} -eq 1 ]] && LoadZfsKey
    fi
}

# Mounts your rootfs
MountRoot()
{
    Info "Mounting rootfs..."

    # Using "" for the ${options} below so that if the user doesn't have any
    # options, the variable ends up expanding back to empty quotes and allows
    # the mount command to keep going.
    PrepForSafeZfsCommand && mount -t zfs -o zfsutil,"${_options}" "${_root}" "${_new_root}" || \
        Fail "Failed to import your zfs root dataset!"

    if [[ ${_refresh} -eq 1 ]] || [[ ! -f "${_rootfs_cache_file}" ]]; then
        # Installs the cache generated by this initramfs run, to the rootfs.
        InstallZpoolCache
    fi
}

# Mounts the /usr directory into your rootfs if it is separate
MountUsrIfNeeded()
{
    local usrFailMessage="Failed to mount \"${_usr}\" onto your rootfs!"

    if [[ ! -z ${_usr} ]]; then
        Info "Mounting /usr onto your rootfs..."
        local targetUsrDirectory="${_new_root}"/usr
        PrepForSafeZfsCommand && mount -t zfs -o zfsutil,"${_options}" "${_usr}" "${targetUsrDirectory}" || Fail "${usrFailMessage}"
    fi
}

# Switches into your root device
SwitchToNewRoot()
{
    Info "Switching into rootfs..." && NewLine
    exec switch_root "${_new_root}" "${_init}"
}

# Runs a zfs specific command that will allow a subsequent zfs mount or unmount
# command to work properly without receiving a "device or resource busy" message.
PrepForSafeZfsCommand()
{
    # Using zfs list primarily for performance (zpool import is way slower
    # because it tries to scan the devices in order to see if it can find
    # anything to import). Normal commands that cause some sort of delay
    # like if [[ $? -eq 0 ]] or 'sleep' will improve the success rate of
    # a subsequent mount command to work, but it still has a high chance to
    # fail. Running a zfs specific command before calling a mount that mounts
    # a zfs dataset seems to have a 100% mount success rate.

    # Using "-H -t -o -s" flags so that if a system has a lot of snapshots,
    # it doesn't slow things down.
    zfs list -H -t filesystem -o name -s name 2>&1 > /dev/null
}

# Installs the zpool.cache to the rootfs
InstallZpoolCache()
{
    Info "Installing zpool.cache into \"${_pool_name}\"..."
    cp -f "${_cache_file}" "${_rootfs_cache_file}"
}

# Single User Mode
SingleUser()
{
    Warn "Booting into single user mode..." && NewLine

    mount --rbind /proc "${_new_root}"/proc
    mount --rbind /dev "${_new_root}"/dev
    mount --rbind /sys "${_new_root}"/sys
    mount --rbind /run "${_new_root}"/run

    _rhostn="rootfs"

    setsid cttyhack /bin/bash -c "chroot ${_new_root} /bin/bash -c 'hostname ${_rhostn}' && chroot ${_new_root} /bin/bash -l"

    # Lazy unmount these devices from the rootfs since they will be fully
    # unmounted from the initramfs environment right after this function
    # is over.
    umount -l "${_new_root}"/proc "${_new_root}"/dev "${_new_root}"/sys "${_new_root}"/run
}

### Utility Functions ###

# Used for displaying information
Info()
{
    echo -e "\e[1;32m[*]\e[0;m ${*}"
}

# Used for input (questions, retrieving feedback from user)
Ask()
{
    echo -en "\e[1;37m[*]\e[0;m ${*}"
}

# Used for warnings
Warn()
{
    echo -e "\e[1;33m[!]\e[0;m ${*}"
}

# Used for flags
Flag()
{
    echo -e "\e[1;34m[+]\e[0;m ${*}"
}

# Used for errors
Fail()
{
    echo -e "\e[1;31m[#]\e[0;m ${*}" && RescueShell
}

# Prints empty line
NewLine()
{
    echo ""
}

# Shows the welcome message
WelcomeMessage()
{
    Info "Welcome to Bliss! [${_version}]"
}

# Prevent kernel from printing on screen
PreventVerboseKernel()
{
    echo 0 > /proc/sys/kernel/printk
}

### Let the games begin ###

LoadVersion
WelcomeMessage
LoadKeymapIfNeeded
MountRequiredDevices || Fail "Failed to mount kernel devices"
PreventVerboseKernel
ParseKernelParameters
LoadModules || Fail "Failed to load kernel modules"
StartUdev

[[ ${_recover} -eq 1 ]] && RescueShell

ZfsTrigger
MountRoot
MountUsrIfNeeded

[[ ${_su} -eq 1 ]] && SingleUser

# Clean up and switch into rootfs
StopUdev
UnmountRequiredDevices || Fail "Failed to unmount kernel devices"
PremountRunOnNewRoot
SwitchToNewRoot || Fail "Failed to switch into your root filesystem"
