# vim: filetype=sh

# The functions in this file are grouped in the following sections:
# - basic routines
# - failsafe mode handling
# - verbosity reduction
# - partitioning and formatting
# - system compression
# - UEFI support
# - package management

# basic routines
# --------------
usage_and_exit()
{
    echo "Usage: debootstick [options] <fs_tree_dir> <out_image_path>" >&2
    exit $1
}

describe_os_support()
{
    cat << EOF
This version of debootstick was tested successfully with the following kinds of
chroot-environments:
- Debian 9 (stretch) as of september 1, 2016
- Debian 8 (jessie)
- Debian 7 (wheezy)
- Ubuntu 16.04 LTS (xenial)
- Ubuntu 14.04 LTS (trusty)
Their binary architecture must be either i386 or amd64.
EOF
}

missing_or_empty()
{
    f="$1"
    if [ ! -e "$f" ]
    then
        echo 1
        return
    else
        if [ "$(cat "$f" | wc -l)" -eq 0 ]
        then
            echo 1
            return
        fi
    fi
    echo 0
}

abspath()
{
    echo "$(cd "$(dirname "$1")" && pwd)/$(basename "$1")"
}

print_last_word()
{
    awk '{print $NF}'
}

busybox_mount()
{
    $busybox_path mount $*
}

# note: when running 'umount <dir>', umount tries to
# find which mount this is about, reading /etc/mtab
# if available, or /proc/mounts.
# Also, some versions of busybox require /proc/self/exe
# to be available otherwise its internal applets,
# including 'mount', will not be found (unless you
# specify "busybox mount" instead of just "mount").
# As a result, /proc should be mounted first (and
# unmounted last, but this is a side effect of
# the failsafe mode handling anyway).
# The caller should mount /proc as needed, and then
# call the function below.
failsafe_mount_sys_and_dev()
{
    failsafe mount -t devtmpfs none /dev
    failsafe mount -t devpts none /dev/pts
    failsafe mount -t sysfs none /sys
}

# divide %1/%2 rounding up
ceil()
{
    divide=$1
    by=$2
    echo $(((divide+by-1)/by))
}

estimated_size_xb()
{
    du -s$1 "$2" | awk '{print $1}'
}

estimated_size_kb()
{
    estimated_size_xb k "$1"
}

estimated_size_mb()
{
    estimated_size_xb m "$1"
}

real_size_human_readable()
{
    du -sh --apparent-size "$1" | awk '{print $1}'
}

device_size_kb()
{
    echo $((
        $(blockdev --getsz $1) /2
    ))
}

update_fstab()
{
    final_root_device="/dev/$1/ROOT"

    # remove the line with "UNCONFIGURED"
    # and add or update the line mounting /
    filtered_content=$(cat /etc/fstab |                 \
                        happy_grep -vw "UNCONFIGURED" | \
                        happy_grep -vw "/"
    )
    cat > /etc/fstab << EOF
$final_root_device / ext4 errors=remount-ro 0 1
$filtered_content
EOF
}

generate_hosts_file()
{
    cat > /etc/hosts << EOF
127.0.0.1   localhost
::1     ip6-localhost ip6-loopback
fe00::0 ip6-localnet
ff00::0 ip6-mcastprefix
ff02::1 ip6-allnodes
ff02::2 ip6-allrouters
EOF
}

# let grub find our virtual device
update_grup_device_map()
{
    loop_device=$1
    cd /
    mkdir -p boot/grub
    cat > boot/grub/device.map << END_MAP
(hd0) $loop_device
END_MAP
    cd - >/dev/null
}

# grep returns non-zero if no line is found.
# in some cases, having no line is expected
# (e.g. in the case of error lines)
# thus the or-true construct.
happy_grep()
{
    grep "$@" || true
}

# get total size needed taking into account
# an overhead.
# $1: the initial size (not taking the overhead into account)
# $2: percent of overhead
# return value: the size needed (such that applying the
#               overhead would get $1 again)
apply_overhead_percent()
{
    echo $((($1)*100/(100-($2))))
}

# apply several overheads
apply_overheads_percent()
{
    size=$1; shift
    while [ ! -z "$1" ]
    do
        size=$(apply_overhead_percent $size "$1")
        shift
    done
    echo $size
}

# check that the given directory looks like an OS
# file structure that we can handle correctly.
check_fs_hierarchy()
{
    fs_tree="$1"
    ls "$fs_tree/bin/echo" >/dev/null 2>&1 || {
        echo "E: No /bin/echo found in $1."
        echo "E: This does not seem to be a chroot environment."
        return 1
    }
    chroot "$fs_tree" echo -n >/dev/null 2>&1 || {
        echo "E: Unable to execute binaries (/bin/echo at least) in the chroot environment."
        echo "E: Please verify:"
        echo "E: - file permissions in the chroot environment"
        echo "E: - that your host CPU is able to run binaries of the target architecture"
        echo "E: Run 'debootstick --help-os-support' for info about compatible environments."
        return 1
    }
    chroot "$fs_tree" which apt-get >/dev/null 2>&1 || {
        echo "E: No apt-get found in $1."
        echo "E: debootstick cannot handle this kind of chroot environment."
        echo "E: Run 'debootstick --help-os-support' for more info."
        return 1
    }
}

check_target_arch()
{
    fs_tree="$1"
    arch="$(chroot "$fs_tree" dpkg --print-architecture 2>&1 || true)"
    case "$arch" in
    "amd64"|"i386") # ok
        ;;
    *)
        echo "E: Got '$arch' while trying to check the architecture of the chroot environment." >&2
        echo "E: debootstick currently cannot handle such a chroot environment." >&2
        echo "E: Run 'debootstick --help-os-support' for more info." >&2
        return 1
        ;;
    esac
    echo "$arch"
}

finalize_fs()
{
    fs_tree="$(cd "$1"; pwd)"
    cd "$fs_tree"

    # clean up
    rm -rf proc/* sys/* dev/* tmp/* \
            $(find run -type f) var/cache/* var/lock

    # move the existing init
    mv sbin/init sbin/init.orig
    cd sbin
    ln -s /opt/debootstick/live/init/first-init.sh init
}

# failsafe mode handling
# ----------------------
# we want to leave the system in a clean state,
# whatever happens.
# for example, if a "disk full" error happens
# in the middle of the chrooted-customization
# step, we should be able to umount what have
# been mounted in the chroot, exit the chroot,
# umount things and remove peripherals created
# by debootstick outside the chroot, before
# exiting.
# we handle this by trapping the EXIT
# of the scripts. Also, each command creating
# a persistent artefact (mounts, devices, etc.)
# is recorded, in order to be able to 'undo the
# command' (i.e. remove the artefacts) if needed.

undo_all()
{
    # run saved failsafe commands prefixed with 'undo_'
    eval "$(echo -n "$failsafe_commands" | \
            awk '{print "undo_" $0}')"

    # flush variable 'failsafe_commands'
    failsafe_commands=""
}

on_sigint()
{
    trap - INT EXIT
    on_exit --from-signal $*
    kill -INT $$
}

on_exit()
{   # save exit code
    res=$?

    # get args
    toplevel=0
    fromsignal=0
    if [ "$1" = "--from-signal" ]
    then
        fromsignal=1
        shift
    fi
    if [ "$1" = "--toplevel" ]
    then
        toplevel=1
        shift
    fi
    cleanup_function=$1

    warn_unexpected_issue=0
    if [ $toplevel -eq 1 ]
    then
        if [ $fromsignal -eq 1 -o $res -gt 0 ]
        then
            warn_unexpected_issue=1
        fi
    fi

    # undo operations (remove artefacts)
    if [ $warn_unexpected_issue -eq 1 ]
    then
        echo
        if [ $fromsignal -eq 1 ]
        then    # signal
            echo "Interrupted."
        else    # error
            echo "E: an error occured."
            echo "E: did you try 'debootstick --help-os-support'?"
        fi
        echo -n "I: restoring a clean state... "
        undo_all
        echo "done"
    else
        undo_all
    fi

    # call an additional cleanup function
    # if provided.
    if [ ! -z "$cleanup_function" ]
    then
        $cleanup_function $res
    fi

    return $res
}

start_failsafe_mode()
{
    # stop if an error occurs
    set -e
    # clean remaining artefacts before exitting
    trap "on_exit $*" EXIT
    trap "on_sigint $*" INT

    # allow with constructs (see f_with function)
    alias with="while f_with"

    # bash does not expand aliases by default,
    # when running a script.
    # busybox sh does, and has no such configuration
    # option (thus the error ignoring construct)
    shopt -s expand_aliases 2>/dev/null || true
}

undo_mount_with_prefix()
{
    # I know 2 usual things that could cause umount
    # to fail with an error reporting that 'device is busy'.
    # Either one process has its current directory on this
    # mount, or there is cached data that was not yet
    # written to disk. We handle these below.
    for last; do true; done # retrieve last arg
    cd / # just in case we would be on the mountpoint
    # some say that a sync request is treated asynchronously.
    # but if a second one comes in, then the first one is
    # forced. Thus the 2 requests in row:
    sync; sync
    $1 umount "$last"
    # try to return to previous dir if possible
    cd - >/dev/null 2>&1 || true
}

undo_mount()
{
    undo_mount_with_prefix "" $*
}

undo_busybox_mount()
{
    undo_mount_with_prefix "$busybox_path" $*
}

undo_mkdir()
{
    for last; do true; done # retrieve last arg
    rm -rf "$last"
}

undo_losetup()
{   # we assume the failsafe command was
    # $ failsafe losetup <loop_device> <file>
    losetup -d "$1"
}

undo_kpartx()
{   # we assume the failsafe command was
    # $ failsafe kpartx -a <disk_device>
    disk_device="$2"

    # we have to detach lvm devices associated
    # to the <disk_device>, then keep the related
    # partition in a busy state otherwise.
    # Retrieving these devices is not so easy...
    partitions=$(kpartx -l $disk_device | \
                    awk '{ print "/dev/mapper/"$1 }')
    vg_names=$(pvs -o vg_name --noheadings $partitions 2>/dev/null || true)
    if [ ! -z "$vg_names" ]
    then
        lv_devices=$(lvs -o vg_name,lv_name --noheadings $vg_names | \
                        awk '{print "/dev/" $1 "/" $2}')
        for lv_device in $lv_devices
        do
            dmsetup remove $lv_device
        done
    fi

    # we can now request the kernel to remove
    # <disk_device> partitions
    kpartx -d "$disk_device"
}

undo_chroot()
{
    exit
}

failsafe()
{
    $*  &&  \
    failsafe_commands="$(
        echo "$*"
        echo -n "$failsafe_commands"
    )"
}

undo()
{
    # undo-ing one failsafe operation only

    # we have to remove this operation from
    # variable 'failsafe_commands'.
    # first, we escape it in order to use
    # it in a sed statement below.
    escaped_cmd="$(
        echo "$*" | \
            sed -e 's/[\/&]/\\&/g')"
    # and now we remove it
    failsafe_commands="$(
        echo -n "$failsafe_commands" | \
            sed -e "/^$escaped_cmd\$/d"
    )"

    # of course we really undo it
    eval "undo_$*"
}

# the function f_with() allows constructs such as:
#
# with mount [...]; do
#   [...]
# done
#
# The unmount-ing will be done at the end of the
# block regardless of what happens inside (issue raised
# or not).
#
# 'with' is actually an alias involving this function
# and a while loop:
# with -> while f_with   (see 'start_failsafe_mode')
#
# we ensure that the while loop stops at the 2nd
# iteration.
f_with()
{
    # save the command
    cmd=$*
    # we need an id to recognise this construct
    with_id=$(echo $cmd | md5sum | awk '{print $1}')
    # let's load the stack of ids we have
    set -- $with_ids_stack

    # if this is a new id...
    if [ "$1" != "$with_id" ]
    then
        # this is a new 'with' construct
        # perform the command requested
        failsafe $cmd
        # update the stack
        with_ids_stack="$with_id $with_ids_stack"
        return 0    # continue the while loop
    else
        # second (and last) time through this 'with' construct
        # pop this id from the stack
        shift; with_ids_stack=$*
        # revert the command
        undo $cmd
        return 1    # stop the while loop
    fi
}


# verbosity reduction
# -------------------
quiet_grub_install()
{
    device=$1

    update_grup_device_map $device

    # grub-install & update-grub print messages to standard
    # error stream although most of these are just
    # informational (or minor bugs). Let's discard them.
    {
        grub-install $device    && \
        update-initramfs -u
        update-grub
        return_code=$?
    } 2>&1 |    happy_grep -v "No error"          | \
                happy_grep -v "Installing"        | \
                happy_grep -v "Generating"        | \
                happy_grep -v "Found .* image:"   | \
                happy_grep -v "lvmetad"           | \
                happy_grep -v "etc.modprobe.d"    | \
                happy_grep -v "^done$" 1>&2

    # the return value we want is the one we stored
    # earlier:
    return $return_code
}

quiet()
{
    $* >/dev/null
}

# partitioning and formatting
# ---------------------------

# Selection of ext4 features:
# We must select only features available on the *oldest* system
# version we want to support. We can get these features by
# creating a sample filesystem on such a system:

# $ cd /tmp
# $ dd of=test.ext4 bs=1G count=0 seek=50
# $ mkfs.ext4 -F -q -L ROOT -T default -m 2 test.ext4
# $ dumpe2fs test.ext4 | grep features

# Note about the '-T default' option:
# We try to create a USB stick as small as possible.
# However, the embedded system may later by copied on a
# potentially large disk.
# Therefore, we specify the option '-T default' to mkfs.ext4.
# This allows to select 'default' ext4 features even if the
# filesystem might be considered 'small' at first.
# This may seem cosmetic but it's not: if initialized with
# '-T small' (or with no -T option and run on a small disk),
# when we move to a large disk, resize2fs apparently
# enables the 'meta_bg' option (supposedly trying to adapt as much
# as possible this 'small' filesystem to a much larger device).
# Since this option is not handled by grub, it prevents the
# system from booting properly.

EXT4_FEATURES=$(cat << EOF
has_journal ext_attr resize_inode dir_index filetype extent
flex_bg sparse_super large_file huge_file uninit_bg dir_nlink
extra_isize
EOF
)

make_root_fs()
{
    features="$(echo $EXT4_FEATURES | tr ' ' ',')"
    mkfs.ext4 -F -q -L ROOT -O "none,$features" -m 2 $1
}

wait_for_device()
{
	while [ ! -e "$1" ]
	do
		sleep 0.1
	done
}

partition_stick()
{
    device=$1
    efi_partition_size_kb=$2
    quiet sgdisk \
            -n 1:0:+${efi_partition_size_kb}K -t 1:ef00 \
            -n 2:0:+${BIOSBOOT_PARTITION_SIZE_KB}K -t 2:ef02 \
            -n 3:0:0 -t 3:8e00 $device
}

create_formatted_image()
{
    image_name=$1
    stick_size_kb=$2
    efi_partition_size_kb=$3
    lvm_vg=$4
    image_file="$5"
    if [ -z "$image_file" ]
    then
        mkdir -p $DBSTCK_TMPDIR/$image_name
        image_file=$DBSTCK_TMPDIR/$image_name/file
    fi

    # create image file
    rm -f "$image_file"
    $DD bs=1024 seek=$stick_size_kb count=0 of="$image_file"

    # partition it
    partition_stick "$image_file" $efi_partition_size_kb

    # let the kernel know about this device
    image_device=$(losetup -f)
    failsafe losetup $image_device "$image_file"
    failsafe kpartx -a $image_device

    # retrieve the partition devices
    set -- $(kpartx -l $image_device | awk '{ print "/dev/mapper/"$1 }')
    image_efipart_device=$1
    image_lvmpart_device=$3

    # create an lvm volume
    wait_for_device $image_lvmpart_device
    quiet pvcreate $image_lvmpart_device
    quiet vgcreate $lvm_vg $image_lvmpart_device
    quiet lvcreate -n ROOT -l 100%FREE $lvm_vg

    # format this lvm volume
    image_lvroot_device=/dev/$lvm_vg/ROOT
    make_root_fs $image_lvroot_device

    # format the efi partition
    wait_for_device $image_efipart_device
    quiet mkfs.vfat -n DBSTCK_EFI $image_efipart_device

    # mount efi partition and root volume
    image_lvroot_mountpoint=$DBSTCK_TMPDIR/$image_name/lvroot
    image_efipart_mountpoint=$DBSTCK_TMPDIR/$image_name/efipart
    mkdir -p $image_lvroot_mountpoint $image_efipart_mountpoint
    failsafe mount $image_lvroot_device $image_lvroot_mountpoint
    failsafe mount $image_efipart_device $image_efipart_mountpoint

    # let the calling code know what we have done
    eval "$(cat << EOF
${image_name}_file="$image_file"
${image_name}_device=$image_device
${image_name}_efipart_device=$image_efipart_device
${image_name}_efipart_mountpoint=$image_efipart_mountpoint
${image_name}_lvm_vg=$lvm_vg
${image_name}_lvmpart_device=$image_lvmpart_device
${image_name}_lvroot_device=$image_lvroot_device
${image_name}_lvroot_mountpoint=$image_lvroot_mountpoint
EOF
    )"
}

release_image()
{
    # read variables with prefix $1=draft or $1=final
    eval "$(cat << EOF
image_name=${1}
image_file="\${${1}_file}"
image_device=\${${1}_device}
image_efipart_device=\${${1}_efipart_device}
image_efipart_mountpoint=\${${1}_efipart_mountpoint}
image_lvm_vg=\${${1}_lvm_vg}
image_lvroot_device=\${${1}_lvroot_device}
image_lvroot_mountpoint=\${${1}_lvroot_mountpoint}
EOF
    )"

    # detach things
    undo mount $image_efipart_device $image_efipart_mountpoint
    undo mount $image_lvroot_device $image_lvroot_mountpoint
    undo kpartx -a $image_device
    undo losetup $image_device "$image_file"
}

# lvm commands should not rely on external daemons in the
# startup scripts
tune_lvm()
{
    mv /etc/lvm/lvm.conf /etc/lvm/lvm.conf.saved
    sed -e 's/\(use_.*d =\).*/\1 0/g' \
        -e 's/udev_sync =.*/udev_sync = 0/g' \
        /etc/lvm/lvm.conf.saved > /etc/lvm/lvm.conf
}

# UEFI support
# ------------

# we generate a standalone UEFI binary, used
# for booting on UEFI systems.
# actually, since we cannot install both grub-pc
# and grub-efi on the embedded system (conflict),
# we install grub-pc only. However, this binary
# image just look for the configuration file generated
# by the grub-pc installation and loads it. Thus
# this UEFI configuration also stays up-to-date.
build_uefi_binary_image()
{
    grub_arch="$1"
    img_name="$2"
    out_binary_path="$(abspath "$img_name")"
    mkdir -p $DBSTCK_TMPDIR/efi/boot/grub
    cd $DBSTCK_TMPDIR/efi
    cat > boot/grub/grub.cfg << EOF
insmod part_gpt
insmod lvm
insmod efi_gop
insmod efi_uga
search --set rootfs --label ROOT
configfile (\$rootfs)/boot/grub/grub.cfg
EOF
    grub-mkstandalone \
            --directory="/usr/lib/grub/$grub_arch/" --format="$grub_arch"   \
            --compress="gz" --output="$out_binary_path"            \
            "boot/grub/grub.cfg"
    cd - >/dev/null # return to previous dir
}

# Package management
# ------------------
package_is_installed()
{
    dpkg-query -W --showformat='${Status}\n' \
                    $1 2>/dev/null | happy_grep -c "^i"
}

list_available_packages()
{
    apt-cache search "$1" | awk '{print $1}'
}

# with some OS versions, package installation
# causes many things to be printed to stderr. Some of those things
# are just informational, others are very minor warnings.
# we will silence them with grep.
INSTALL_GREP_PATTERN="$(cat << EOF | tr -d '\n'
(delaying package configuration)|(Done.)|(^Moving old)|(^Running)|
(^update-initramfs: deferring)|(^Examining)|(^run-parts: executing)|
(^update-initramfs: Generating)|(^initrd.img)|(points to)|
(doing nothing)|(^vmlinu)|(connect to Upstart)|(policy-rc.d denied)|
(Creating config)|(^$)|(start and stop actions)|(^Created symlink)|
(Initializing machine ID)|(:$)|(etc.modprobe.d)
EOF
)"

install_packages()
{
    packages=$*

    # disable service startup at package installation
    if [ -e /usr/sbin/policy-rc.d ]
    then
	mv /usr/sbin/policy-rc.d /usr/sbin/policy-rc.d.saved
    fi
    echo exit 101 > /usr/sbin/policy-rc.d
    chmod +x /usr/sbin/policy-rc.d

    apt-get -qq --no-install-recommends -o=Dpkg::Use-Pty=0 \
		install $packages 2>&1 >/dev/null | \
        happy_grep -vE "$INSTALL_GREP_PATTERN" 1>&2

    # restore policy-rc.d conf
    if [ -e /usr/sbin/policy-rc.d.saved ]
    then
	mv /usr/sbin/policy-rc.d.saved /usr/sbin/policy-rc.d
    else
        rm /usr/sbin/policy-rc.d
    fi
}
