#!/bin/bash
#
# Copyright 2012 Hewlett-Packard Development Company, L.P.
#
# 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.

set -eE

# Prevent perl from complaining a lot, but also remove any unexpected side-effects
# of $LANG varying between build hosts
export LANG=C

# Store our initial environment and command line args for later
export DIB_ARGS="$@"
export DIB_ENV=$(export | grep ' DIB_.*=')

SCRIPTNAME=$(basename $0)
SCRIPT_HOME=$(dirname $(readlink -f $0))
if [ -d $SCRIPT_HOME/../share/diskimage-builder ]; then
    export _PREFIX=$SCRIPT_HOME/../share/diskimage-builder
else
    export _PREFIX=$SCRIPT_HOME/..
fi
export _LIB=$_PREFIX/lib
source $_LIB/die

IS_RAMDISK=0
if [ "$SCRIPTNAME" == "ramdisk-image-create" ]; then
  IS_RAMDISK=1
fi

function show_options () {
    echo "Usage: ${SCRIPTNAME} [OPTION]... [ELEMENT]..."
    echo
    echo "Options:"
    echo "    -a i386|amd64|armhf -- set the architecture of the image(default amd64)"
    echo "    -o imagename -- set the imagename of the output image file(default image)"
    echo "    -t qcow2,tar -- set the image types of the output image files (default qcow2)"
    echo "       File types should be comma separated"
    echo "    -x -- turn on tracing"
    echo "    -u -- uncompressed; do not compress the image - larger but faster"
    echo "    -c -- clear environment before starting work"
    echo "    --image-size size -- image size in GB for the created image"
    echo "    --image-cache directory -- location for cached images(default ~/.cache/image-create)"
    echo "    --max-online-resize size -- max number of filesystem blocks to support when resizing."
    echo "       Useful if you want a really large root partition when the image is deployed."
    echo "       Using a very large value may run into a known bug in resize2fs."
    echo "       Setting the value to 274877906944 will get you a 1PB root file system."
    echo "       Making this value unnecessarily large will consume extra disk space "
    echo "       on the root partition with extra file system inodes."
    echo "    --min-tmpfs size -- minimum size in GB needed in tmpfs to build the image"
    echo "    --no-tmpfs -- do not use tmpfs to speed image build"
    echo "    --offline -- do not update cached resources"
    echo "    --qemu-img-options -- option flags to be passed directly to qemu-img."
    echo "       Options need to be comma separated, and follow the key=value pattern."
    echo "    --root-label label -- label for the root filesystem.  Defaults to 'cloudimg-rootfs'."
    echo "    --ramdisk-element -- specify the main element to be used for building ramdisks."
    if [ "$IS_RAMDISK" == "0" ]; then
        echo "    -n skip the default inclusion of the 'base' element"
        echo "    -p package[,package,package] -- list of packages to install in the image"
    fi
    echo "    -h|--help -- display this help and exit"
    echo
    echo "ELEMENTS_PATH will allow you to specify multiple locations for the elements."
    echo
    echo "NOTE: At least one distribution root element must be specified."
    echo
    echo "Examples:"
    if [ "$IS_RAMDISK" == "0" ]; then
        echo "    ${SCRIPTNAME} -a amd64 -o ubuntu-amd64 vm ubuntu"
        echo "    export ELEMENTS_PATH=~/source/tripleo-image-elements/elements"
        echo "    ${SCRIPTNAME} -a amd64 -o fedora-amd64-heat-cfntools vm fedora heat-cfntools"
    else
        echo "    ${SCRIPTNAME} -a amd64 -o fedora-deploy deploy fedora"
        echo "    ${SCRIPTNAME} -a amd64 -o ubuntu-ramdisk ramdisk ubuntu"
    fi
}

INSTALL_PACKAGES=""
IMAGE_TYPES=("qcow2")
COMPRESS_IMAGE="true"
DIB_ROOT_LABEL=""
TEMP=`getopt -o a:ho:t:xucnp: -l no-tmpfs,offline,help,min-tmpfs:,image-size:,image-cache:,max-online-resize:,qemu-img-options:,ramdisk-element:,root-label: -n $SCRIPTNAME -- "$@"`
if [ $? -ne 0 ] ; then echo "Terminating..." >&2 ; exit 1 ; fi

# Note the quotes around `$TEMP': they are essential!
eval set -- "$TEMP"

while true ; do
    case "$1" in
        -a) export ARCH=$2; shift 2 ;;
        -o) export IMAGE_NAME=$2; shift 2 ;;
        -t) IFS="," read -a IMAGE_TYPES <<< "$2"; export IMAGE_TYPES ; shift 2 ;;
        -h|--help) show_options; exit 0;;
        -x) shift; set -x;;
        -u) shift; export COMPRESS_IMAGE="";;
        -c) shift ; export CLEAR_ENV=1;;
        -n) shift; export SKIP_BASE="1";;
        -p) IFS="," read -a INSTALL_PACKAGES <<< "$2"; export INSTALL_PACKAGES ; shift 2 ;;
        --image-size) export DIB_IMAGE_SIZE=$2; shift 2;;
        --image-cache) export DIB_IMAGE_CACHE=$2; shift 2;;
        --max-online-resize) export MAX_ONLINE_RESIZE=$2; shift 2;;
        --min-tmpfs) export DIB_MIN_TMPFS=$2; shift 2;;
        --no-tmpfs) shift; export DIB_NO_TMPFS=1;;
        --offline) shift; export DIB_OFFLINE=1;;
        --qemu-img-options) QEMU_IMG_OPTIONS=$2; shift 2;;
        --root-label) export DIB_ROOT_LABEL=$2; shift 2;;
        --ramdisk-element) RAMDISK_ELEMENT=$2; shift 2;;
        --) shift ; break ;;
        *) echo "Internal error!" ; exit 1 ;;
    esac
done

export DIB_IMAGE_CACHE=${DIB_IMAGE_CACHE:-~/.cache/image-create}
mkdir -p $DIB_IMAGE_CACHE

if [ "$CLEAR_ENV" = "1" -a "$HOME" != "" ]; then
  echo "Re-execing to clear environment."
  echo "(note this will prevent much of the local_config element from working)"
  exec -c $0 "$@"
fi

source $_LIB/img-defaults
source $_LIB/common-functions
source $_LIB/img-functions

if [ "$IS_RAMDISK" == "1" ]; then
  source $_LIB/ramdisk-defaults
  source $_LIB/ramdisk-functions
fi

arg_to_elements "$@"

if [ "${#IMAGE_TYPES[@]}" = "1" ]; then
  export IMAGE_NAME=${IMAGE_NAME%%\.${IMAGE_TYPES[0]}}
fi

# NOTE: Tuning the rootfs uuid works only for ext filesystems.
# Rely on the below environment variable only for ext filesystems.
export DIB_IMAGE_ROOT_FS_UUID=$(uuidgen -r)

# FS_TYPE isn't available until after we source img-defaults
if [ -z "$DIB_ROOT_LABEL" ]; then
    # NOTE(bnemec): XFS has a limit of 12 characters for filesystem labels
    # Not changing the default for other filesystems to maintain backwards compatibility
    if [ "$FS_TYPE" = "xfs" ]; then
        DIB_ROOT_LABEL="img-rootfs"
    else
        DIB_ROOT_LABEL="cloudimg-rootfs"
    fi
fi

mk_build_dir
create_base
# This variable needs to be propagated into the chroot
echo "export DIB_ROOT_LABEL=\"${DIB_ROOT_LABEL}\"" > $TMP_HOOKS_PATH/environment.d/10-dib-root-label.bash
run_d extra-data
# Run pre-install scripts. These do things that prepare the chroot for package installs
run_d_in_target pre-install
do_extra_package_install
# Call install scripts to pull in the software users want.
run_d_in_target install
run_d_in_target post-install
# ensure we do not have a lost+found directory in the root folder
# that could cause copy to fail (it will be created again later
# when creating the file system, if it needs such directory)
if [ -e "$TMP_BUILD_DIR/mnt/lost+found" ]; then
    sudo rm -rf "$TMP_BUILD_DIR/mnt/lost+found"
fi
# Free up /mnt
unmount_image
mv $TMP_BUILD_DIR/mnt $TMP_BUILD_DIR/built

MKFS_OPTS=""

if [ -n "$DIB_IMAGE_SIZE" ]; then
  truncate -s${DIB_IMAGE_SIZE}G $TMP_IMAGE_PATH
else
  # in kb*0.60 - underreport to get a slightly bigger device
  # Rounding down size so that is is a multiple of 64, works around a bug in
  # qemu-img that may occur when compressing raw images that aren't a multiple
  # of 64k. https://bugs.launchpad.net/ubuntu/+source/linux/+bug/1180021
  _NEEDED_SIZE=$(sudo du --block-size=600 -x -s ${TMP_BUILD_DIR}/built | \
                 awk ' { print $1  - ( $1 % 64) } ')
  truncate -s${_NEEDED_SIZE}K $TMP_IMAGE_PATH
  if [ "$FS_TYPE" = "ext4" ] ; then
    # Very conservative to handle images being resized a lot
    # Without -J option specified, default journal size will be set to 32M
    # and online resize will be failed with error of needs too many credits.
    MKFS_OPTS="-i 4096 -J size=64"
  fi
fi

if [ -n "$MAX_ONLINE_RESIZE" ]; then
    MKFS_OPTS="$MKFS_OPTS -E resize=$MAX_ONLINE_RESIZE"
fi

LOOPDEV=$(sudo losetup --show -f $TMP_IMAGE_PATH)
export EXTRA_UNMOUNT="detach_loopback $LOOPDEV"
export IMAGE_BLOCK_DEVICE=$LOOPDEV
eval_run_d block-device "IMAGE_BLOCK_DEVICE="
sudo mkfs $MKFS_OPTS -t $FS_TYPE -L ${DIB_ROOT_LABEL} ${IMAGE_BLOCK_DEVICE}
# Tuning the rootfs uuid works only for ext filesystems.
if echo "$FS_TYPE" | grep -q "^ext"; then
    sudo tune2fs ${IMAGE_BLOCK_DEVICE} -U ${DIB_IMAGE_ROOT_FS_UUID}
fi
mkdir $TMP_BUILD_DIR/mnt
sudo mount ${IMAGE_BLOCK_DEVICE} $TMP_BUILD_DIR/mnt
sudo mv -t $TMP_BUILD_DIR/mnt ${TMP_BUILD_DIR}/built/*
mount_proc_dev_sys
run_d_in_target finalise
finalise_base

for X in ${!IMAGE_TYPES[@]} ; do
  if [ "${IMAGE_TYPES[$X]}" == "tar" ]; then
    sudo tar -C ${TMP_BUILD_DIR}/mnt -cf $IMAGE_NAME.tar --exclude ./sys \
             --exclude ./proc --xattrs --xattrs-include=\* .
    sudo chown $USER: $IMAGE_NAME.tar
    unset IMAGE_TYPES[$X]
  fi
done

unmount_image

if [ "$IS_RAMDISK" == "0" ]; then
  for IMAGE_TYPE in ${IMAGE_TYPES[@]} ; do
    compress_and_save_image $IMAGE_NAME.$IMAGE_TYPE
  done
  rm $TMP_IMAGE_PATH
  cleanup_dirs
else
  # This is a ramdisk build, we have already extracted the kernel and ramdisk
  # by this point.
  rm $TMP_IMAGE_PATH
fi
# All done!
trap EXIT