#!/bin/sh -e
#
# Cumulus Linux Make Binary Image
#
# Copyright 2015 Cumulus Networks, Inc.
# All rights reserved.
#

CL_MKIMAGE_QUIET=""
CL_MKIMAGE_DRYRUN=""
CL_MKIMAGE_DEBUG=""
CL_MKIMAGE_CONTROL=""
CL_MKIMAGE_APTARCHIVE=""
CL_MKIMAGE_SYSROOTTAR=""
CL_MKIMAGE_EI="${CL_MKIMAGE_EI:-/usr/share/cumulus-onie/embedded-installer.tgz}"
CL_MKIMAGE_SELFEXTRACTOR_TPL="${CL_MKIMAGE_SELFEXTRACTOR_TPL:-/usr/share/cumulus-onie/bin-exe.template}"
CL_MKIMAGE_ARCH="${CL_MKIMAGE_ARCH:-$(dpkg --print-architecture)}"
CL_MKIMAGE_PRESEED=""
CL_MKIMAGE_DTBDIR=""
CL_MKIMAGE_CWD="$(pwd)"

logerr()
{
	echo "E: $@" 1>&2
}

loginfo()
{
	test -n "${CL_MKIMAGE_QUIET}" || echo "I: $@"
}

logdbg()
{
	test "xyes" = "x${CL_MKIMAGE_DEBUG}" && echo "D: $@" || true
}

run_cmd()
{
	local pre=""
	test -n "$CL_MKIMAGE_DRYRUN" && pre="Would Run"
	test -n "$CL_MKIMAGE_DRYRUN" -o -n "$CL_MKIMAGE_DEBUG" &&
		loginfo "${pre} Command: $@"
	test -n "$CL_MKIMAGE_DRYRUN" && return 0
	eval "$@" || {
		logerr "$@"
		return 1
	}
	return 0
}

require_superuser()
{
	test "$(id -u)" = "0" -o -n "$CL_MKIMAGE_DRYRUN" || {
		logerr "You need superuser privileges for this operation."
		return 1
	}
	return 0
}

cleanup()
{
	test -n "${CL_MKIMAGE_DEBUG}" && return 0
	test -n "${RAMDISKDIR}" && umount -R "${RAMDISKDIR}" >/dev/null 2>&1 || true
	test -n "${RAMDISKDIR}" && rm -rf "${RAMDISKDIR}"
}

# TODO: add an option --dtbdir as some architectures need a dtb to boot
usage()
{
	echo "Usage: onie-mkimage [OPTIONS] IMAGENAME KERNEL INITRD"
	test "$1" = "long" || return;
	echo "
Make an ONIE compatiable executable install binary.

Options:
  -q,--quiet           suppress progress messages
  -n,--dry-run         pretend to do actions, don't actually do them
  -d,--debug           enable extra debigging messages as well as
                       keep temporary directories
  -h,--help            this help output
  -v,--version         installer version information
  -a,--arch=ARCH       architecture to build image for (default: ${CL_MKIMAGE_ARCH})
  --control=FILE       RFC-8222 style image control file to use
  --apt=FILE           Tarball of an APT archive containing the base system to
                       install. Can be compressed with gzip, bzip2, or xz
  --dtb=DTBDIR         dtb directory containing all the dtbs used by various
                       platforms that should be supported by this installer
  --sysroot=DIR|FILE   Tarball of a sysroot which represents the base system
                       to install. Can be compressed with gzip, bzip2, or xz
  -p,--preseed=FILE    provide a preseed file to use during the install

IMAGENAME
    The name the executable tarball will be named once the image creating
    process is complete. Image will be stored in the current working
    directory.

KERNEL
    Cumulus-Installer kernel

INITRD
    Cumulus-Installer initrd
"
}


ARGS=$(getopt -o "dqnhva:p:" -l "debug,quiet,dry-run,help,version,arch:,control:,apt:,sysroot:,preseed:,dtb:" -- "$@")

if test $? -ne 0; then
	exit 1
fi

eval set -- "$ARGS"

while true; do
	case "$1" in
	-d|--debug)
		CL_MKIMAGE_DEBUG="yes"
		shift;
		;;
	-q|--quiet)
		CL_MKIMAGE_QUIET="yes"
		shift;
		;;
	-h|--help)
		usage "long"
		exit 0;
		shift;
		;;
	-v|--version)
		echo "onie-mkimage UNKNOWN"
		echo ""
		echo "Copyright 2015 Cumulus Networks, Inc."
		echo "All rights reserved."
		exit 0;
		shift;
		;;
	-n|--dry-run)
		export CL_MKIMAGE_DRYRUN="yes"
		shift;
		;;
	-a|--arch)
		shift;
		CL_MKIMAGE_ARCH="$1"
		shift;;
	--control)
		shift;
		CL_MKIMAGE_CONTROL="$1"
		if test ! -e "${CL_MKIMAGE_CONTROL}"; then
			logerr "--control ${CL_MKIMAGE_CONTROL} does not exist"
			exit 2
		fi
		shift;;
	--apt)
		shift;
		CL_MKIMAGE_APTARCHIVE="$1"
		if test ! -e "${CL_MKIMAGE_APTARCHIVE}"; then
			logerr "--apt ${CL_MKIMAGE_APTARCHIVE} does not exist"
			exit 2
		fi
		shift;;
	--sysroot)
		shift;
		CL_MKIMAGE_SYSROOTTAR="$1"
		if test ! -e "${CL_MKIMAGE_SYSROOTTAR}"; then
			logerr "--sysroot ${CL_MKIMAGE_SYSROOTTAR} does not exist"
			exit 2
		fi
		shift;;
	-p|--preseed)
		shift;
		CL_MKIMAGE_PRESEED="$1"
		shift;;
	--dtb)
		shift;
		CL_MKIMAGE_DTBDIR="$1"
		shift;;
	--)
		shift;
		break;
		;;
	esac
done

test $# -lt 3 && {
	logerr "Usage: you must specify IMAGENAME, KERNEL, and INITRD"
	usage
	exit 2
}

IMAGENAME="$1"
shift
if test -e "${IMAGENAME}"; then
	logerr "IMAGENAME already exists: ${IMAGENAME}"
	exit 2
fi

KERNEL="$1"
shift
if test -z "${KERNEL}" || test ! -e "${KERNEL}"; then
	logerr "Invalid KERNEL provided: ${KERNEL}"
	exit 2
fi

INITRD="$1"
shift
if test -z "${INITRD}" || test ! -e "${INITRD}"; then
	logerr "Invalid INITRD provided: ${INITRD}"
	exit 2
fi

trap cleanup EXIT
trap cleanup TERM
trap cleanup QUIT
trap cleanup ABRT
trap cleanup HUP
trap cleanup INT

RAMDISKDIR=$(mktemp -p "${CL_MKIMAGE_CWD}" -d ramdisk.d.XXXXXXXX)
logdbg "ramdisk tempdir: ${RAMDISKDIR}"

PAYLOADDIR=$(mktemp -p "${RAMDISKDIR}" -d payload.d.XXXXXXXX)
logdbg "payload tempdir: ${PAYLOADDIR}"

PAYLOADTAR=$(mktemp -p "${RAMDISKDIR}" payload.tar.XXXXXXX)
logdbg "payload tar: ${PAYLOADTAR}"

if test -n "${CL_MKIMAGE_SYSROOTTAR}"; then
	loginfo "Copy sysroot archive to payload directory..."
	SYSROOTTAR="${PAYLOADDIR}/sysroot.tar"
	run_cmd cp "${CL_MKIMAGE_SYSROOTTAR}" "${SYSROOTTAR}" || {
		logerr "Problem copying sysroot archive to payload directoy"
		exit 1
	}
fi

if test -n "${CL_MKIMAGE_APTARCHIVE}"; then
	loginfo "Copy APT archive to payload directory..."
	APTARCHIVE="${PAYLOADDIR}/repo-archive.tar"
	run_cmd cp "${CL_MKIMAGE_APTARCHIVE}" "${APTARCHIVE}" || {
		logerr "Problem copying apt archive to payload directoy"
		exit 1
	}
fi

loginfo "Copy KERNEL and INITRD to payload directory..."
run_cmd cp "$KERNEL" "${PAYLOADDIR}/kernel" || {
	logerr "Problem copying KERNEL to payload directoy"
	exit 1
}
KERNEL="${PAYLOADDIR}/kernel"

run_cmd cp "$INITRD" "${PAYLOADDIR}/initrd" || {
	logerr "Problem copying INITRD to payload directoy"
	exit 1
}
INITRD="${PAYLOADDIR}/initrd"

if test -n "${CL_MKIMAGE_PRESEED}"; then
	loginfo "Modifying INITRD with custom preseed file..."
	run_cmd cp "${CL_MKIMAGE_PRESEED}" "${PAYLOADDIR}/preseed.cfg" || {
		logerr "Problem copying PRESEED file"
		exit 1
	}

	(cd "${PAYLOADDIR}" && \
	 ls "preseed.cfg" | cpio --quiet -o -H newc >preseed.cpio && \
	 gzip preseed.cpio) || {
		logerr "Problem creating cpio archive for preseed file"
		exit 1
	}

	cat "${INITRD}" "${PAYLOADDIR}/preseed.cpio.gz" >"${PAYLOADDIR}/initrd.tmp" || {
		logerr "problem concatenating preseed and initrd cpio archives"
		exit 1
	}

	run_cmd mv "${PAYLOADDIR}/initrd.tmp" "${INITRD}" || {
		logerr "problem renaming temporary initrd"
		exit 1
	}

	run_cmd rm ${PAYLOADDIR}/preseed* || {
		logerr "problem removing temporary preseed files"
		exit 1
	}
fi

if test "armel" = "${CL_MKIMAGE_ARCH}" -o "powerpc" = "${CL_MKIMAGE_ARCH}"; then
	loginfo "Modifying INITRD wrapping in uImage..."
	case "${CL_MKIMAGE_ARCH}" in
	armel)
		arch="arm";;
	powerpc)
		arch="powerpc";;
	*)
		logerr "trying to uImage wrap an unknown architecture"
		exit 1
	esac

	if test -z "${CL_MKIMAGE_DRYRUN}"; then
		mkimage -A "$arch" -O linux -T ramdisk -C none -a 0x0 \
			-n "cumulus-installer initramfs" \
			-d "$INITRD" \
			"${INITRD}.tmp" 1>/dev/null || exit 1
		mv "${INITRD}.tmp" "$INITRD" || exit 1
	fi
fi

if test -n "${CL_MKIMAGE_DTBDIR}"; then
	loginfo "Adding dtbs to payload..."
	run_cmd mkdir -p "${PAYLOADDIR}/dtbs"
	for dtb in "${CL_MKIMAGE_DTBDIR}/*.dtb"; do
		run_cmd cp "${dtb}" "${PAYLOADDIR}/dtbs/"
	done
fi

loginfo "Calculate minimum disk size required..."
# We have a hard minimum size limit of 1024MB just in case a sysroot is not
# provided and is not larger than this minimum.
minsize=$((1024 * 1024 * 1024))
logdbg "Min disk size: $minsize bytes"

loginfo "Calculate minimum RAM required..."
# We calculate the minimum RAM required as double the total installer size,
# we have a hard minimum of 512MB. Doing double installer size allows us to
# roughly account for having to store the sysroot tarball into ram and then
# still be able to run the installer.
minram=$((512 * 1024 * 1024))
if test -z "${CL_MKIMAGE_DRYRUN}"; then
	installersize=$(stat --format="%s" "$KERNEL")
	installersize=$(($installersize + $(stat --format="%s" "$INITRD")))
	if test -n "${APTARCHIVE}"; then
		installersize=$(($installersize + $(stat --format="%s" \
				"${APTARCHIVE}")))
	fi
	if test -n "${SYSROOTTAR}"; then
		installersize=$(($installersize + $(stat --format="%s" \
				"${SYSROOTTAR}")))
	fi
	installersize=$(($installersize * 2))
	if test "$minram" -lt "$installersize"; then
		minram="$installersize"
	fi
fi
logdbg "Min RAM size: $minram bytes"

loginfo "Extracting embedded-installer..."
run_cmd tar -C "${PAYLOADDIR}" -xzf "${CL_MKIMAGE_EI}" || {
	logerr "Problems extracting embedded-installer; $CL_MKIMAGE_EI"
	exit 1
}

loginfo "Generating control file..."
# Generate RFC-8222 style control file.  Keep these names distinct so
# we can easily "grep" for them.  Use dashes to separate words.
if test -n "${CL_MKIMAGE_CONTROL}"; then
	loginfo "Applying user supplied image control file..."
	run_cmd cp "${CL_MKIMAGE_CONTROL}" $PAYLOADDIR/control
else
	# inspect os-release to pickup release variables
	os_release="${SYSROOTDIR}/etc/os-release"
	release=$(grep VERSION= $os_release | sed -r 's/VERSION="(.*)"/\1/')
	desc=$(grep PRETTY_NAME= $os_release | sed -r 's/PRETTY_NAME="(.*)"/\1/')
	cat <<EOF > $PAYLOADDIR/control
Description: $desc
OS-Release: $release
Architecture: ${CL_MKIMAGE_ARCH}
Switch-Architecture: bcm-${CL_MKIMAGE_ARCH}
Date: $(date -R)
Homepage: http://www.cumulusnetworks.com/
EOF
fi

# Append additional info from onie-mkimage
cat <<EOF >> $PAYLOADDIR/control
Min-Disk-Size: ${minsize}
Min-Ram-Size: ${minram}
mkimage-version: UNKNOWN
EOF

loginfo "Creating payload.tar..."
# No need to compress here, as the inputs are already compressed.
run_cmd tar -C "${PAYLOADDIR}" -cf "${PAYLOADTAR}" . || {
	logerr "Problems creating payload.tar archive."
	exit 1
}

loginfo "Calculate sha1sum of payload.tar..."
payloadsha1=$(cat "${PAYLOADTAR}" | sha1sum | awk '{print $1}')
if test -z "${payloadsha1}"; then
	logerr "Problem generating sha1sum for playload.tar."
	exit 1
fi

loginfo "Fixup self-extractor.sh template and glue to front of tarball..."
if test -z "${CL_MKIMAGE_DRYRUN}"; then
	mkdir -p "$(dirname ${IMAGENAME})"
	sed -e "s/@IMAGE_SHA1@/${payloadsha1}/" \
		"${CL_MKIMAGE_SELFEXTRACTOR_TPL}" >${IMAGENAME}
	cat "${PAYLOADTAR}" >>${IMAGENAME}
fi

loginfo "chmod tarball to executable..."
run_cmd chmod +x "${IMAGENAME}"

loginfo "Calculate sha1sum of install binary..."
if test -z "${CL_MKIMAGE_DRYRUN}"; then
	c=$(pwd)
	base=$(basename ${IMAGENAME})
	cd "$(dirname ${IMAGENAME})"
	sha1sum "${base}" > ${base}.sha1
	cd $c
fi
loginfo "Success:  Install image is ready at ${IMAGENAME}"
exit 0
