#! /bin/sh

# $NetBSD: mkimage,v 1.3.2.3 2013/01/23 00:04:28 yamt 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}"
}

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

# Return the filesystem size for an ls -l or tar -xvf list
# Directories and symlinks in tar are 0 size, we assume one block
# (which is too much), we round up by the fragment size the rest.
getfssize() {
	local bsize="$1"
	local fsize="$2"

	awk -v fsize=${fsize} -v bsize=${bsize} '
	NF >= 9 && $1 != "tar:" {
		if ($5 == 0)
			tot += bsize;
		else
			tot += ((int)(($5 + fsize - 1) / fsize)) * fsize;
	}
	END {
		printf("%d\n", tot);
	}'
}

# find the size of the gzipped files in a .tgz archive
# Directories appear as 0, so count them as one block
# and round up files to a fragment.
sizeone() {
	if [ ! -f "$1" ]
	then
		echo "$PROG: Missing set $1" 1>&2
		echo 0
		return;
	fi

        case "$1" in 
        *.tgz|*.tar.gz|*.tbz|*.tar.bz2|*.txz|*.tar.xz)
                tar tvzf "$1" | getfssize ${bsize} ${fsize}
                ;;
        *)
                echo 0
                ;; 
        esac
}


# Return the usable filesystem size in bytes, given the total size in bytes,
# and optionally block and fragment sizes
getffssize() {
	local bytes="$1"
	local barg
	local farg
	local overhead

	if [ -n "$2" ]
	then
		barg="-b $2"
		if [ -n "$3" ]
		then
			farg="-f $3"
		fi
	fi

	overhead=$(newfs -N ${farg} ${barg} -s "${bytes}b" -F /dev/null |
	    awk '/using/ {
		printf("%d\n", substr($6, 1, length($6) - 3) * 1024 * 1024);
	    }'
	)
	echo $(( ${bytes} - ${overhead} ))
}

# Compute the size of an ffs filesystem that can fit x bytes.
# Instead of duplicating the newfs calculations here we let
# it do the job, using binary search.
makeffssize() {
	local bytes=$1
	local bsize=$2
	local fsize=$3
	local max=$(( 2 * ${bytes} ))
	local min="${bytes}"
	local cur
	local res
	while true; do
		cur="$(( ( ${max} + ${min} ) / 2 ))"
		res="$(getffssize "${cur}" ${bsize} ${fsize})"
#		echo ${min} ${cur} ${max} ${res} ${bytes} 1>&2
		if [ "${res}" -eq "${bytes}" ]
		then
		    break
		elif [ "$(( ${min} + 1 ))" -ge "${max}" ]
		then
		    break
		elif [ "${res}" -lt "${bytes}" ]
		then
		    min="${cur}"
		elif [ "${res}" -gt "${bytes}" ]
		then
		    max="${cur}"
		fi
	done
	echo "${cur}"
}

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


DIR="$(dirname "$0")"
PROG="$(basename "$0")"
bar="==="
sudo=
mnt="${TMPDIR:-/tmp}/image.$$"
src="/usr/src"
obj="/usr/obj"

sets="base comp etc games man misc modules text"
xsets="xbase xcomp xetc xfont xserver"

# Presumable block and fragment size.
bsize=16384
fsize=2048
mtob=$(( 1024 * 1024 ))

# 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)	sets="$sets $xsets";;
	*)	usage;;
	esac
done

trap finish 0 1 2 3 15

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

# calculate the set bytes
setbytes=0
echo -n "${bar} computing set sizes ("
b=
for s in ${sets}; do
	one="$(sizeone ${setsdir}/${s}.tgz)"
	echo -n "$b$s=$(( ${one} / ${mtob} ))MB"
	setbytes=$(( ${setbytes} +  ${one} ))
	b=" "
done
echo "): $(( ${setbytes} / ${mtob} ))MB ${bar}"

# calculate size of custom files
custbytes=0
if [ -d "${custom}" ]; then
	custbytes=$(ls -lR "${custom}" | getfssize ${bsize} ${fsize})
fi
echo "${bar} computing custom sizes: $(( ${custbytes} / ${mtob} ))MB ${bar}"

# how many bytes
rawbytes="$(( ${setbytes} + ${custbytes} ))"
echo -n "${bar} computing ffs filesystem size for $(( ${rawbytes} / ${mtob} ))MB: "
ffsbytes="$(makeffssize "${rawbytes}")"
ffsmb=$(( ${ffsbytes} / ${mtob} ))
echo " ${ffsmb}MB ${bar}"

# total in MB
total=$(( ${ffsmb} + ${overhead} ))
echo "${bar} overhead: ${overhead}MB ${bar}"

if [ $size -eq 0 ]; then
        # auto-size the pkgs fs
        newsize=${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
	newsize=${size}
fi

echo "${bar} making a new ${newsize} MB image in ${image} ${bar}"
dd if=/dev/zero of=${image} bs=1m count=${newsize} 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 -n "${bar} installing sets:"
(cd ${mnt} &&
	for s in ${sets}; do
		ss="${setsdir}/${s}.tgz"
		if [ -f "${ss}" ]; then
			echo -n " ${s}"
			${sudo} tar xpzf "${ss}"
		fi
	done
)
echo " ${bar}"

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