#! /bin/sh

# $NetBSD: mkimage,v 1.5 2013/01/13 20:58:38 christos Exp $

# Copyright (c) 2012 Alistair Crooks <agc@NetBSD.org>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
# NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
# THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

# find next available vnd, from kre
next_avail ()
{
	local dev="$1"
	local N=$(( ${#dev} + 1 ))
	local unit units

	units=$(
		sysctl -n hw.disknames		|
			tr ' ' '\012'		|
			grep '^'"${dev}"'[0-9]'	|
			sort -n -k 1.$N			)

	test -z "${units}" && {
		test -e "/dev/${dev}0a" || {
			echo >&2 "No ${dev}s available!"
			return 1
		}
		echo "${dev}0"
		return
	}

	N=0
	for unit in ${units}
	do
		if [ "${unit}" = "${dev}${N}" ]
		then
			N=$(( N + 1 ))
		else
			echo "${dev}${N}"
			return
		fi
	done

	test -e /dev/"${dev}${N}a" || {
		echo >&2 "All ${dev}s in use"
		return 1
	}

	echo "${dev}${N}"
}

# find the size of the gzipped files in a .tgz archive
sizeone() {
	if [ ! -f "$1" ]
	then
		echo "$PROG: Missing set $1" 1>&2
		echo 0
		return;
	fi
        case "$1" in 
        *.tgz|*.tar.gz)
                tar tvzf "$1" | awk '{ tot += $5 } END { print tot }'
                ;;
        *.tbz|*.tar.bz2)
                tar tvjf "$1" | awk '{ tot += $5 } END { print tot }' 
                ;;
        *)
                echo 0
                ;; 
        esac
}

usage() {
	cat << EOF 1>&2
Usage: $PROG [-S <setsdir>] [-c <custom-files-dir>] [-h <host-arch>] [-s <size>]
EOF
exit 1
}

finish() {
    cleanup
    ${sudo} umount ${mnt}
    ${sudo} vnconfig -u ${vnddev}
}

trap finish 0 1 2 3 15
DIR="$(dirname "$0")"
PROG="$(basename "$0")"
bar="==="
sudo=
mnt="/tmp/image.$$"
src="/usr/src"
obj="/usr/obj"

# First pass for options to get the host
OPTS="S:c:h:s:x"
while getopts "$OPTS" f
do
	case $f in
	h)	h="$OPTARG";;
	*)	;;
	esac
done

if [ -z "$h" ]
then
	usage
fi

if [ ! -f "${DIR}/conf/${h}.conf" ]
then
	echo $PROG: ${DIR}/conf/${h}.conf is not present 1>&2
	exit 1
fi

. "${DIR}/conf/${h}.conf"

OPTIND=1
while getopts "$OPTS" f
do
	case $f in
	S)	setsdir="$OPTARG";;
	c)	custom="$OPTARG";;
	h)	;;
	s)	size="$OPTARG";;
	x)	set -x;;
	*)	;;
	esac
done

shift $(( "$OPTIND" - 1 ))
if [ -n "$1" ]; then
	# take the next argument as being the image name
	image="$1"
	shift
fi

total=0
for s in ${sets}; do
	one="$(sizeone ${setsdir}/${s}.tgz)"
	total=$(( ${total} +  ${one} ))
done
# calculate size of custom files
custsize=0
if [ -d "${custom}" ]; then
	custsize=$(ls -lR "${custom}" | awk 'NF == 9 { tot += $5 } END { print tot }')
fi
total=$(( ( ( ${total} + ${custsize} ) / 1000000 ) + ${overhead} ))
if [ $size -eq 0 ]; then
        # auto-size the pkgs fs
        size=${total}
else
        # check that we've been given enough space
        if [ ${total} -gt ${size} ]; then
                echo "$PROG: Given size is ${size} MB, but we need ${total} MB" >&2
                exit 1
        fi
fi

echo "${bar} making a new ${size} MB image in ${image} ${bar}"
dd if=/dev/zero of=${image} bs=1m count=${size} conv=sparse

vnddev=$(next_avail vnd)
echo "${bar} mounting image via vnd ${vnddev} ${bar}"
${sudo} vnconfig ${vnddev} ${image}
${sudo} mkdir -p ${mnt}
make_filesystems

${sudo} mkdir -p ${mnt}/etc ${mnt}/dev

echo "${bar} installing sets ${bar}"
(cd ${mnt} &&
	for s in ${sets}; do
		if [ -f "${s}" ]; then
			echo ${s}
			${sudo} tar xpzf ${setsdir}/${s}.tgz
		fi
	done
)

echo "${bar} performing customisations ${bar}"

make_fstab

${sudo} cat > ${mnt}/etc/rc.conf << EOF
#
# see rc.conf(5) for more information.
#
# Use program=YES to enable program, NO to disable it. program_flags are
# passed to the program on the command line.
#

# Load the defaults in from /etc/defaults/rc.conf (if it's readable).
# These can be overridden below.
#
if [ -r /etc/defaults/rc.conf ]; then
        . /etc/defaults/rc.conf
fi

# If this is not set to YES, the system will drop into single-user mode.
#
rc_configured=YES

hostname=${h}

EOF

customize

for d in ${specialdirs}; do
	${sudo} mkdir -p ${mnt}/${d}
done

if [ \( -n "${custom}" \) -a \( -d "${custom}" \) ]; then
	echo "${bar} user customisations from files in ${custom} ${bar}"
	(cd ${custom} && ${sudo} pax -rwpe . ${mnt})
fi

exit 0