# Copyright 2012 Hewlett-Packard Development Company, L.P.
# All Rights Reserved.
#
# 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.

function unmount_image () {
    # Calling sync before helps ensure the mount isn't busy when you unmount it.
    # Previously observing having disk corruption issues; one possibility is
    # qemu-nbd not flushing dirty pages on disconnect?
    # https://bugs.launchpad.net/diskimage-builder/+bug/1214388
    sync

    # unmount from the chroot
    # Don't use TMP_MOUNT_PATH here, it might not have been set.
    unmount_dir "$TMP_BUILD_DIR/mnt"
    if [ -n "$EXTRA_UNMOUNT" ]; then
        $EXTRA_UNMOUNT
    fi
}

function trap_cleanup() {
    exitval=$?
    cleanup
    exit $exitval
}

function cleanup () {
    unmount_image
    cleanup_dirs
}

function ensure_nbd () {
    NBD=`which qemu-nbd` || true
    if [ -z "$NBD" ]; then
        echo "qemu-nbd is not found in your PATH"
        echo "Please install it on your system"
        exit 1
    fi
    # prep nbd for mounting
    (lsmod | grep '^nbd ' >/dev/null) || sudo modprobe nbd max_part=16
}

function ensure_sudo () {
    sudo echo "Ensuring sudo is available"
}

# Helper function to run a command inside the chroot
function run_in_target () {
    # Force the inclusion of /usr/local/bin in PATH, this is needed for some
    # distros that does not include /usr/local/bin in the sudoers secure_path.
    # Note that we're not expanding PATH during argument processing, the \$
    # will preserve the PATH syntax until after the sh command runs
    cmd="PATH=\$PATH:/usr/local/bin ; $@"
    # -E to preserve http_proxy
    ORIG_HOME=$HOME
    export HOME=/root
    sudo -E chroot $TMP_MOUNT_PATH sh -c "$cmd"
    export HOME=$ORIG_HOME
}

# Helper function to run a directory of scripts inside the chroot
function run_d_in_target () {
    check_element
    # If we can find a directory of hooks to run in the target filesystem, bind
    # mount it into the target and then execute run-parts in a chroot
    if [ -d ${TMP_HOOKS_PATH}/$1.d ] ; then
      sudo mkdir $TMP_MOUNT_PATH/tmp/in_target.d
      sudo mount --bind ${TMP_HOOKS_PATH} $TMP_MOUNT_PATH/tmp/in_target.d
      sudo mount -o remount,ro,bind ${TMP_HOOKS_PATH} $TMP_MOUNT_PATH/tmp/in_target.d
      check_break before-$1 run_in_target bash
      trap "check_break after-error run_in_target bash" ERR
      run_in_target dib-run-parts /tmp/in_target.d/$1.d
      trap - ERR
      check_break after-$1 run_in_target bash
      sudo umount -f $TMP_MOUNT_PATH/tmp/in_target.d
      sudo rmdir $TMP_MOUNT_PATH/tmp/in_target.d
    fi
}

function prepare_first_boot () {
    check_break before-first-boot run_in_target bash
    if [ -d ${TMP_HOOKS_PATH}/first-boot.d ] ; then
      sudo cp -t $TMP_MOUNT_PATH/etc/ -a $TMP_HOOKS_PATH/first-boot.d
    fi
    check_break after-first-boot run_in_target bash
}

function finalise_base () {
    TARGET_ROOT=$TMP_MOUNT_PATH run_d cleanup
    # If the file has been set immutable, we probably want to keep it
    if lsattr $TMP_MOUNT_PATH/etc/resolv.conf | grep '^....i' >/dev/null ; then
        # We're keeping the contents of resolv.conf set in the elements,
        # so remove the old saved file
        sudo rm -f $TMP_MOUNT_PATH/etc/resolv.conf.ORIG
    else
        # Remove the resolv.conf we created above
        sudo rm -f $TMP_MOUNT_PATH/etc/resolv.conf
        # Move the original back
        if [ -L $TMP_MOUNT_PATH/etc/resolv.conf.ORIG ] || [ -f $TMP_MOUNT_PATH/etc/resolv.conf.ORIG ] ; then
            sudo mv $TMP_MOUNT_PATH/etc/resolv.conf.ORIG $TMP_MOUNT_PATH/etc/resolv.conf
        fi
    fi
}

function compress_and_save_image () {
    # Recreate our image to throw away unnecessary data
    test $IMAGE_TYPE != qcow2 && COMPRESS_IMAGE=""
    if [ -n "$QEMU_IMG_OPTIONS" ]; then
        EXTRA_OPTIONS="-o $QEMU_IMG_OPTIONS"
    else
        EXTRA_OPTIONS=""
    fi
    qemu-img convert ${COMPRESS_IMAGE:+-c} -f raw $TMP_IMAGE_PATH -O $IMAGE_TYPE $EXTRA_OPTIONS $1-new

    OUT_IMAGE_PATH=$1-new
    finish_image $1
}

function do_extra_package_install () {
  # Install any packages that were requested with the -p command line option
  if [ "$INSTALL_PACKAGES" != "" ]; then
    run_in_target install-packages ${INSTALL_PACKAGES[@]}
  fi
}

function copy_elements_lib () {
  sudo mkdir -p $TMP_MOUNT_PATH/lib/diskimage-builder
  sudo cp -t $TMP_MOUNT_PATH/lib/diskimage-builder $_LIB/elements-functions
}

# Dig up the initrd and kernel.
function select_boot_kernel_initrd () {
    TARGET_ROOT=$1
    BOOTDIR=$TARGET_ROOT/boot
    if [ -n "${DIB_BAREMETAL_KERNEL_PATTERN:-}" -a -n "${DIB_BAREMETAL_INITRD_PATTERN:-}" ]; then
        KERNEL=$(basename $(eval ls -1rv "$BOOTDIR/${DIB_BAREMETAL_KERNEL_PATTERN}" | head -1))
        RAMDISK=$(basename $(eval ls -1rv "$BOOTDIR/${DIB_BAREMETAL_INITRD_PATTERN}" | head -1))
    elif [ -f $TARGET_ROOT/etc/redhat-release ]; then

        # Prioritize PAE if present
        KERNEL=$(ls -1rv $BOOTDIR/vmlinuz* | grep PAE | grep -v debug | head -1 || echo "")
        KERNEL=${KERNEL:-$(ls -1rv $BOOTDIR/vmlinuz* | grep -v debug | head -1 || echo "")}
        if [ ! $KERNEL ]; then
           echo "No suitable kernel found."
           exit 1
        fi

        KERNEL=$(basename $KERNEL)
        KERNEL_VERSION=${KERNEL#vmlinuz-}
        RAMDISK=$(basename $(ls $BOOTDIR/initramfs-$KERNEL_VERSION.img) || echo "")
        if [ ! $RAMDISK ]; then
            echo "Can't find an initramfs for the $KERNEL_VERSION version of the kernel."
            exit 1
        fi
    elif [ -f $TARGET_ROOT/etc/debian_version ]; then
        KERNEL=$(basename $(ls -1rv $BOOTDIR/vmlinuz*generic 2>/dev/null || ls -1rv $BOOTDIR/vmlinuz* | head -1))
        RAMDISK=$(basename $(ls -1rv $BOOTDIR/initrd*generic 2>/dev/null || ls -1rv $BOOTDIR/initrd* | head -1))
    elif [ -f $TARGET_ROOT/etc/SuSE-release ]; then
        KERNEL=$(basename $(readlink -e $BOOTDIR/vmlinuz))
        RAMDISK=$(basename $(readlink -e $BOOTDIR/initrd))
    else
        echo "ERROR: Unable to detect operating system"
        exit 1
    fi
}