#!/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

# Set/override locale.  This ensures consistency in sorting etc.  We
# need to choose a lowest-common denominator locale, as this is
# applied when running in the building chroot too (maybe a bug and we
# should prune this?).  Thus "C" --centOS 7 doesn't include C.utf-8
# (fedora does, centos 8 probably will).  Note: LC_ALL to really
# override this; it overrides LANG and all other LC_ vars
export LC_ALL=C

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

SCRIPTNAME=$(basename $0)

if [ -z "$_LIB" ]; then
    echo "_LIB not set!"
    exit 1
fi

_BASE_ELEMENT_DIR=$(${DIB_PYTHON_EXEC:-python} -c '
import diskimage_builder.paths
diskimage_builder.paths.show_path("elements")')

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|arm64 -- 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,tgz,squashfs,vhd,docker,aci,raw -- set the image types of the output image files (default qcow2)"
    echo "       File types should be comma separated. VHD outputting requires the vhd-util"
    echo "       executable be in your PATH. ACI outputting requires the ACI_MANIFEST "
    echo "       environment variable be a path to a manifest file."
    echo "    -x -- turn on tracing (use -x -x for very detailed tracing)."
    echo "    -u -- uncompressed; do not compress the image - larger but faster"
    echo "    -c -- clear environment before starting work"
    echo "    --logfile -- save run output to given logfile"
    echo "    --checksum -- generate MD5 and SHA256 checksum files for the created image"
    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 "    --mkfs-options -- option flags to be passed directly to mkfs."
    echo "       Options should be passed as a single string value."
    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."
    echo "       Defaults to 'ramdisk'.  Should be set to 'dracut-ramdisk' for platforms such"
    echo "       as RHEL and CentOS that do not package busybox."
    echo "    --install-type -- specify the default installation type. Defaults to 'source'. Set to 'package' to use package based installations by default."
    echo "    --docker-target -- specify the repo and tag to use if the output type is docker. Defaults to the value of output imagename"
    if [ "$IS_RAMDISK" == "0" ]; then
        echo "    -n skip the default inclusion of the 'base' element"
        echo "    -p package[,p2...] [-p p3] -- extra packages to install in the image.  Runs once, after 'install.d' phase.  Can be specified mulitple times"
    fi
    echo "    -h|--help -- display this help and exit"
    echo "    --version -- display version and exit"
    echo
    echo "Environment Variables:"
    echo "  (this is not a complete list)"
    echo
    echo "  * ELEMENTS_PATH: specify external locations for the elements.  As for \$PATH"
    echo "  * DIB_NO_TIMESTAMP: no timestamp prefix on output.  Useful if capturing output"
    echo "  * DIB_QUIET: do not output log output to stdout.  See --logfile"
    echo
    echo "NOTE: At least one distribution root element must be specified."
    echo
    echo "NOTE: If using the VHD output format you need to have a patched version of vhd-util installed for the image"
    echo "      to be bootable. The patch is available here: https://github.com/emonty/vhd-util/blob/master/debian/patches/citrix"
    echo "      and a PPA with the patched tool is available here: https://launchpad.net/~openstack-ci-core/+archive/ubuntu/vhd-util"
    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
}

function show_version() {
    ${DIB_PYTHON_EXEC:-python} -c "from diskimage_builder import version; print(version.version_info.version_string())"
}

DIB_DEBUG_TRACE=${DIB_DEBUG_TRACE:-0}
INSTALL_PACKAGES=""
IMAGE_TYPES=("qcow2")
COMPRESS_IMAGE="true"
ROOT_LABEL=""
DIB_DEFAULT_INSTALLTYPE=${DIB_DEFAULT_INSTALLTYPE:-"source"}
MKFS_OPTS=""
ACI_MANIFEST=${ACI_MANIFEST:-}
DOCKER_TARGET=""
LOGFILE=""
TEMP=`getopt -o a:ho:t:xucnp: -l checksum,no-tmpfs,offline,help,version,min-tmpfs:,image-size:,image-cache:,max-online-resize:,mkfs-options:,qemu-img-options:,ramdisk-element:,root-label:,install-type:,docker-target:,logfile: -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;;
        --version) show_version; exit 0;;
        -x) shift; DIB_DEBUG_TRACE=$(( $DIB_DEBUG_TRACE + 1 ));;
        -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=( ${INSTALL_PACKAGES[@]} ${_INSTALL_PACKAGES[@]} ) ; shift 2 ;;
        --checksum) shift; export DIB_CHECKSUM=1;;
        --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;;
        --mkfs-options) MKFS_OPTS=$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) ROOT_LABEL=$2; shift 2;;
        --ramdisk-element) RAMDISK_ELEMENT=$2; shift 2;;
        --install-type) DIB_DEFAULT_INSTALLTYPE=$2; shift 2;;
        --docker-target) export DOCKER_TARGET=$2; shift 2 ;;
        --logfile) export LOGFILE=$2; shift 2 ;;
        --) shift ; break ;;
        *) echo "Internal error!" ; exit 1 ;;
    esac
done

export DIB_DEBUG_TRACE

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

# We send stdout & stderr through "outfilter" which does timestamping,
# basic filtering and log file output.
_TS_FLAG=""
if [[ "${DIB_NO_TIMESTAMP:-0}" -eq 1 ]]; then
    _TS_FLAG="--no-timestamp"
fi
_QUIET_FLAG="-v"
if [[ "${DIB_QUIET:-0}" -eq 1 ]]; then
    _QUIET_FLAG=""
fi
_LOGFILE_FLAG=""
if [[ -n "${LOGFILE}" ]]; then
    echo "Output logs going to: ${LOGFILE}"
    _LOGFILE_FLAG="-o ${LOGFILE}"
fi
exec 1> >( ${DIB_PYTHON_EXEC:-python} $_LIB/outfilter.py ${_TS_FLAG} ${_QUIET_FLAG} ${_LOGFILE_FLAG} ) 2>&1


# Display the current file/function/line in the debug output
function _ps4 {
    IFS=" " called=($(caller 0))
    local f=$(readlink -f ${called[2]})
    printf "%-80s " "$f:${called[1]}:${called[0]}"
}
export -f _ps4
export PS4='+ $(_ps4):   '

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

echo "diskimage-builder version $(show_version)"

# If no elements are specified theres no way we can succeed
if [ -z "$*" ]; then
    echo "ERROR: At least one distribution root element must be specified"
    exit 1
fi
arg_to_elements "$@"

# start tracing after most boilerplate
if [ ${DIB_DEBUG_TRACE} -gt 0 ]; then
    set -x
fi

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

# Check for required tools early on
for X in ${!IMAGE_TYPES[@]}; do
    case "${IMAGE_TYPES[$X]}" in
        qcow2)
            if ! type qemu-img > /dev/null 2>&1; then
                echo "qcow2 output format specified but qemu-img executable not found."
                exit 1
            fi
            ;;
        tgz)
            # Force tar to be created.
            IMAGE_TYPES+=('tar')
            ;;
        vhd)
            if ! type vhd-util > /dev/null 2>&1; then
                echo "vhd output format specified but no vhd-util executable found."
                exit 1
            fi
            ;;
        squashfs)
            if ! type mksquashfs > /dev/null 2>&1; then
                echo "squashfs output format specified but no mksquashfs executable found."
                exit 1
            fi
            ;;
        docker)
            if ! type docker > /dev/null 2>&1; then
                echo "docker output format specified but no docker executable found."
                exit 1
            fi
            if [ -z "$DOCKER_TARGET" ]; then
                echo "Please set --docker-target."
                exit 1
            fi
            ;;
    esac
done

# NOTE: fstrim is on most all recent systems. It is provided by the util-linux
# package.
if ! type fstrim > /dev/null 2>&1; then
    echo "fstrim utility is not found. This is provided by util-linux package"
    echo "Please check your PATH variable is set correctly"
    exit 1
fi

# xattr support cannot be relied upon with tmpfs builds
# some kernels supoprt it, some don't
if [[ -n "${GENTOO_PROFILE}" ]]; then
    if [[ "${GENTOO_PROFILE}" =~ "hardened" ]]; then
        echo 'disabling tmpfs for gentoo hardened build'
        export DIB_NO_TMPFS=1
    fi
fi

mk_build_dir

# Create the YAML file with the final and raw configuration for
# the block device layer.
mkdir -p ${TMP_BUILD_DIR}/block-device
BLOCK_DEVICE_CONFIG_YAML=${TMP_BUILD_DIR}/block-device/config.yaml
block_device_create_config_file "${BLOCK_DEVICE_CONFIG_YAML}"

# Write out the parameter file
DIB_BLOCK_DEVICE_PARAMS_YAML=${TMP_BUILD_DIR}/block-device/params.yaml
export DIB_BLOCK_DEVICE_PARAMS_YAML
cat >${DIB_BLOCK_DEVICE_PARAMS_YAML} <<EOF
config: ${BLOCK_DEVICE_CONFIG_YAML}
image-dir: ${TMP_IMAGE_DIR}
root-fs-type: ${FS_TYPE}
root-label: ${ROOT_LABEL}
mount-base: ${TMP_BUILD_DIR}/mnt
build-dir: ${TMP_BUILD_DIR}
EOF

dib-block-device init

# Need to get the real root label because it can be overwritten
# by the BLOCK_DEVICE_CONFIG.
DIB_ROOT_LABEL=$(dib-block-device getval root-label)
export DIB_ROOT_LABEL

# Need to get the real fs type for the root filesystem
DIB_ROOT_FSTYPE=$(dib-block-device getval root-fstype)
export DIB_ROOT_FSTYPE

# retrieve mount points so we can reuse in elements
DIB_MOUNTPOINTS=$(dib-block-device getval mount-points)
export DIB_MOUNTPOINTS

create_base
# This variable needs to be propagated into the chroot
mkdir -p $TMP_HOOKS_PATH/environment.d
echo "export DIB_DEFAULT_INSTALLTYPE=\${DIB_DEFAULT_INSTALLTYPE:-\"${DIB_DEFAULT_INSTALLTYPE}\"}" > $TMP_HOOKS_PATH/environment.d/11-dib-install-type.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
# Call install scripts to pull in the software users want.
run_d_in_target install
do_extra_package_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

# save xtrace state, as we always want to turn it off to avoid
# spamming the logs with du output below.
xtrace=$(set +o | grep xtrace)

# temp file for holding du output
du_output=${TMP_BUILD_DIR}/du_output.tmp

if [ -n "$DIB_IMAGE_SIZE" ]; then
    du_size=$(echo "$DIB_IMAGE_SIZE" | awk '{printf("%d\n",$1 * 1024 *1024)}')
else
    set +o xtrace
    echo "Calculating image size (this may take a minute)..."
    sudo du -a -c -x ${TMP_BUILD_DIR}/built > ${du_output}
    # the last line is the total size from "-c".
    # scale this by 0.6 to create a slightly bigger image
    du_size=$(tail -n1 ${du_output} | cut -f1 | awk '{print int($1 / 0.6)}')
    $xtrace
fi

if [[ "${DIB_SHOW_IMAGE_USAGE:-0}" != 0 ]]; then
    set +o xtrace
    if [ ! -f "$du_output" ]; then
        sudo du -a -c -x ${TMP_BUILD_DIR}/built > ${du_output}
    fi

    du_output_show="sort -nr ${du_output} |
                     numfmt --to=iec-i --padding=7
                       --suffix=B --field=1 --from-unit=1024"

    # by default show the 10MiB and greater files & directories -- a
    # dir with lots of little files will still show up, but this helps
    # signal:noise ratio
    if [[ ${DIB_SHOW_IMAGE_USAGE_FULL:-0} == 0 ]]; then
        # numfmt will start giving a decimal place when < 10MiB
        du_output_show+="| egrep 'MiB|GiB|TiB|PiB' | grep -v '\..MiB'"
        echo "================================="
        echo "Image size report (files > 10MiB)"
        echo "================================="
    else
        echo "================="
        echo "Image size report"
        echo "================="
    fi

    eval ${du_output_show}

    echo
    echo "===== end image size report ====="
    echo

    $xtrace
fi

rm -f ${du_output}

if [ "$DIB_ROOT_FSTYPE" = "ext4" ] ; then
  # Very conservative to handle images being resized a lot
  # We set journal size to 64M so our journal is large enough when we
  # perform an FS resize.
  MKFS_OPTS="-i 4096 -J size=64 $MKFS_OPTS"
  # Grow the image size to account for the journal, only if the user
  # has not asked for a specific size.
  if [ -z "$DIB_IMAGE_SIZE" ]; then
    du_size=$(( $du_size + 65536 ))
  fi
fi

# 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
export DIB_IMAGE_SIZE=$(echo "$du_size" | awk ' { if ($1 % 64 != 0) print $1 + 64 - ( $1 % 64); else print $1; } ')

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

export TMP_IMAGE_DIR
# Try the 'old fashioned' way calling the block device
# phase. If this gives no result, use the configuration based approach:
eval_run_d block-device "IMAGE_BLOCK_DEVICE="

if [ -z ${IMAGE_BLOCK_DEVICE} ] ; then
    # For compatibily reasons in addition to the YAML configuration
    # there is the need to handle the old environment variables.
    echo "image-size: ${DIB_IMAGE_SIZE}KiB" >> ${DIB_BLOCK_DEVICE_PARAMS_YAML}

    if [ -n "${MKFS_OPTS}" ] ; then
        echo "root-fs-opts: '${MKFS_OPTS}'" >> ${DIB_BLOCK_DEVICE_PARAMS_YAML}
    fi

    # After changeing the parameters, there is the need to
    # re-run dib-block-device init because some value might
    # change based on the new set parameters.
    dib-block-device init

    # values to dib-block-device: using the YAML config and
    dib-block-device create

    # This is the device (/dev/loopX).  It's where to install the
    # bootloader.
    IMAGE_BLOCK_DEVICE=$(dib-block-device getval image-block-device)
    export IMAGE_BLOCK_DEVICE

    # Similar to above, but all mounted devices.  This is handy for
    # some bootloaders that have multi-partition layouts and want to
    # copy things to different places other than just
    # IMAGE_BLOCK_DEVICE.  "eval" this into an array as needed
    IMAGE_BLOCK_DEVICES=$(dib-block-device getval image-block-devices)
    export IMAGE_BLOCK_DEVICES

    # Write the fstab
    dib-block-device writefstab
fi

# XXX: needed?
LOOPDEV=${IMAGE_BLOCK_DEVICE}

# At this point, dib-block-device has created the raw image file
# (IMAGE_BLOCK_DEVICE) and mounted all the partitions under
# $TMP_BUILD_DIR/mnt for us.  We can now copy into the final image.

# 'mv' is not usable here - especially when a top level directory
# has the same name as a mount point of a partition.  If so, 'mv'
# will complain:
# mv: inter-device move failed: '...' to '...'; \
#       unable to remove target: Device or resource busy
# therefore a 'cp' and 'rm' approach is used.
sudo cp -ra ${TMP_BUILD_DIR}/built/* $TMP_BUILD_DIR/mnt
sudo rm -fr ${TMP_BUILD_DIR}/built/*

mount_proc_dev_sys
run_d_in_target finalise
finalise_base

for X in ${!IMAGE_TYPES[@]} ; do
  if [[ " tar aci " =~ "${IMAGE_TYPES[$X]}" ]]; then
    if [ "${IMAGE_TYPES[$X]}" = "aci" ]; then
      sudo tar -C ${TMP_BUILD_DIR}/mnt -cf $IMAGE_NAME.aci --exclude ./sys \
               --exclude ./proc --xattrs --xattrs-include=\* \
               --transform 's,^.,rootfs,S' .
      if [ -n "$ACI_MANIFEST" ]; then
        cp $ACI_MANIFEST ${TMP_BUILD_DIR}/manifest
        sudo tar -C ${TMP_BUILD_DIR} --append -f $IMAGE_NAME.aci manifest
      else
        echo "No ACI_MANIFEST specified. An ACI_MANIFEST must be specified for"
        echo " this image to be usable."
      fi
    else
      sudo tar -C ${TMP_BUILD_DIR}/mnt -cf $IMAGE_NAME.tar --exclude ./sys \
               --exclude ./proc --xattrs --xattrs-include=\* .
    fi
    sudo chown $USER: $IMAGE_NAME.${IMAGE_TYPES[$X]}
    unset IMAGE_TYPES[$X]
  elif [ "${IMAGE_TYPES[$x]}" == "squashfs" ]; then
    sudo mksquashfs ${TMP_BUILD_DIR}/mnt $IMAGE_NAME.squash -comp xz \
                    -noappend -root-becomes ${TMP_BUILD_DIR}/mnt \
                    -wildcards -e "proc/*" -e "sys/*" -no-recovery
  elif [ "${IMAGE_TYPES[$X]}" == "docker" ]; then
    sudo tar -C ${TMP_BUILD_DIR}/mnt -cf - --exclude ./sys \
             --exclude ./proc --xattrs --xattrs-include=\* . \
        | sudo docker import - $DOCKER_TARGET
    unset IMAGE_TYPES[$X]
  fi
done

if [[ ! $IMAGE_ELEMENT =~ no-final-image ]]; then
  # Prep filesystem by discarding all unused space
  fstrim_image
fi

# Unmount and cleanup the /mnt and /build subdirectories, to save
# space before converting the image to some other format.
# XXX ? needed?
export EXTRA_UNMOUNT=""
unmount_image

TMP_IMAGE_PATH=$(dib-block-device getval image-path)
export TMP_IMAGE_PATH

# remove all mounts
dib-block-device umount
dib-block-device cleanup

cleanup_build_dir

if [[ (! $IMAGE_ELEMENT =~ no-final-image) && "$IS_RAMDISK" == "0" ]]; then
  has_raw_type=
  for IMAGE_TYPE in ${IMAGE_TYPES[@]} ; do
    # We have to do raw last because it is destructive
    if [ "$IMAGE_TYPE" = "raw" ]; then
      has_raw_type=1
    elif [ "$IMAGE_TYPE" != "squashfs" ]; then
      compress_and_save_image $IMAGE_NAME.$IMAGE_TYPE
    fi
  done
  if [ -n "$has_raw_type" ]; then
    IMAGE_TYPE="raw"
    compress_and_save_image $IMAGE_NAME.$IMAGE_TYPE
  fi
fi

# Remove the leftovers, i.e. the temporary image directory.
cleanup_image_dir

# All done!
trap EXIT