#!/bin/sh -e
#
# Cumulus Linux Make Binary Image
#
# Copyright 2022 NVIDIA Corporation.
# Copyright 2015,2016,2018 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_FIRMWARE="${CL_MKIMAGE_FIRMWARE:-$(dirname ${CL_MKIMAGE_EI})/firmware}"
CL_MKIMAGE_ARCH="${CL_MKIMAGE_ARCH:-$(dpkg --print-architecture)}"
CL_MKIMAGE_PRESEED=""
CL_MKIMAGE_DTBDIR=""
CL_MKIMAGE_CWD="$(pwd)"
CL_MKIMAGE_KEY_FILE=""
CL_MKIMAGE_CERT_FILE=""

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

logwarn()
{
	echo "W: $@" 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
  --dtb=DIR            dtb directory
  --cert=FILE          signer certificate (must also use --key)
  --key=FILE           signing key (must also use --cert)

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:,cert:,key:" -- "$@")

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;;
	--cert)
		shift;
		CL_MKIMAGE_CERT_FILE="$1"
		shift;;
	--key)
		shift;
		CL_MKIMAGE_KEY_FILE="$1"
		shift;;
	--)
		shift;
		break;
		;;
	esac
done

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

if [ "$CL_MKIMAGE_CERT_FILE" = "" -a "$CL_MKIMAGE_KEY_FILE" != "" ]; then
        logerr "Must specify both certificate and key files for signing."
        usage
        exit 2
fi

if [ "$CL_MKIMAGE_CERT_FILE" != "" -a "$CL_MKIMAGE_KEY_FILE" = "" ]; then
        logerr "Must specify both certificate and key files for signing."
        usage
        exit 2
fi

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"

if test "armel" = "${CL_MKIMAGE_ARCH}" -o "powerpc" = "${CL_MKIMAGE_ARCH}"; then
	loginfo "Modifying KERNEL, wrapping in uImage..."
	case "${CL_MKIMAGE_ARCH}" in
	armel)
		arch="arm"
		load_addr=0x61008000
		entry_addr=0x61008000
		;;
	powerpc)
		arch="powerpc"
		logerr "PowerPC kernel not currently supported"
		exit 1
		;;
	*)
		logerr "trying to uImage wrap an unknown architecture"
		exit 1
	esac

	if test -z "${CL_MKIMAGE_DRYRUN}"; then
		mkimage -A "$arch" -O linux -T kernel -C none \
			-a $load_addr -e $entry_addr \
			-n "cumulus-installer kernel" \
			-d "$KERNEL" \
			"${KERNEL}.tmp" 1>/dev/null || exit 1
		mv "${KERNEL}.tmp" "$KERNEL" || exit 1
	fi
fi

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
}

FIRMWARE_TGZ="${CL_MKIMAGE_FIRMWARE}-${CL_MKIMAGE_ARCH}.tgz"
if test -f "$FIRMWARE_TGZ" ; then
	loginfo "Adding firmware for ${CL_MKIMAGE_ARCH}..."
	run_cmd tar -C "${PAYLOADDIR}" -xzf "${FIRMWARE_TGZ}" || {
		logerr "Problems extracting firmware: $FIRMWARE_TGZ"
		exit 1
	}
fi

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

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

loginfo "Find size of payload ${PAYLOADTAR} ..."
payload_image_size=$(stat -c %s "${PAYLOADTAR}")
if test -z "${payload_image_size}"; then
	logerr "Problem generating size for payload ${PAYLOADTAR}."
	exit 1
fi

loginfo "Fixup ${CL_MKIMAGE_SELFEXTRACTOR_TPL} template and glue to front of payload ${PAYLOADTAR} to make install image ${IMAGENAME}..."
if test -z "${CL_MKIMAGE_DRYRUN}"; then
	mkdir -p "$(dirname ${IMAGENAME})"
	sed -e "s/@IMAGE_SHA256@/${payloadsha256}/" \
	    -e "s/@PAYLOAD_IMAGE_SIZE@/${payload_image_size}/" \
		"${CL_MKIMAGE_SELFEXTRACTOR_TPL}" >${IMAGENAME}
	cat "${PAYLOADTAR}" >>${IMAGENAME}
fi

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

if [ "$CL_MKIMAGE_CERT_FILE" != "" ]; then
        loginfo "Signing ${IMAGENAME} with ${CL_MKIMAGE_CERT_FILE} and ${CL_MKIMAGE_KEY_FILE}"
        CL_MKIMAGE_SIGNATURE_FILE=$( mktemp XXXXXXXX )
        openssl cms -sign -nosmimecap \
                -signer ${CL_MKIMAGE_CERT_FILE} \
                -inkey ${CL_MKIMAGE_KEY_FILE} \
                -binary -in ${IMAGENAME} \
                -outform pem -out ${CL_MKIMAGE_SIGNATURE_FILE}
        cat ${CL_MKIMAGE_SIGNATURE_FILE}
        cat ${CL_MKIMAGE_SIGNATURE_FILE} >> ${IMAGENAME}
        rm ${CL_MKIMAGE_SIGNATURE_FILE}
else
        logwarn "No signing certificate and key specified, ${IMAGENAME} is unsigned"
fi

for sha in 1 256 ; do
	loginfo "Calculate sha${sha}sum of install image ${IMAGENAME}..."
	if test -z "${CL_MKIMAGE_DRYRUN}"; then
		c=$(pwd)
		base=$(basename ${IMAGENAME})
		cd "$(dirname ${IMAGENAME})"
		sha${sha}sum "${base}" > ${base}.sha${sha}
		cd $c
	fi
done

loginfo "Success:  Install image is ready at ${IMAGENAME}"
exit 0
