#!/bin/bash

# Gentoaster VM build tool
# Licensed under GPL v3, see COPYING file
# Usage: ./create_image.sh <configuration file path>
# Prerequisites: extlinux, qemu-img, sfdisk

# Early error handling (we'll override this later)

handle_error() {
	echo $1
	exit 1
}

# Do some sanity checks first
if [ "$(id -u)" != "0" ]; then
   handle_error "Gentoaster must run with root permissions!"
fi
hash qemu-img 2>&- || handle_error "Gentoaster requires qemu-img, but it's not installed."
hash extlinux 2>&- || handle_error "Gentoaster requires extlinux, but it's not installed."
hash sfdisk 2>&- || handle_error "Gentoaster requires sfdisk, but it's not installed."

# Figure out where we are
RUNNING_DIRECTORY=$(cd `dirname $0` && pwd)

# Load shflags to parse flags
. ${RUNNING_DIRECTORY}/lib/shflags

DEFINE_string 'config' '' 'configuration to build from' 
DEFINE_string 'proxy' '' 'HTTP proxy to use for emerges'
DEFINE_boolean 'cachedkernel' false 'use a cached kernel (shortens build time)'
DEFINE_boolean 'compress' false 'compressed the finished image'

FLAGS "$@" || exit 1
eval set -- "${FLAGS_ARGV}"

if [ -z ${FLAGS_config} ]; then
	handle_error "Gentoaster requires a valid configuration to be passed using the --config flag"
fi

# Parse the configuration we took as an arg
source ${RUNNING_DIRECTORY}/parse_config.sh ${FLAGS_config} 2>/dev/null || handle_error "Error parsing build configuration"

# Generate a few helper variables using the configuration file

IMAGE_NAME="${BUILD_ID}.image"
IMAGE_MEGABYTES=$(( ${BOOT_MEGABYTES} + ${SWAP_MEGABYTES} + ${ROOT_MEGABYTES} + 1 ))
IMAGE_BYTES=$(( ${IMAGE_MEGABYTES} * 1024 * 1024 ))
IMAGES_OUTPUT_PATH=`pwd`
IMAGE_WORK_PATH="${IMAGES_OUTPUT_PATH}/${BUILD_ID}"
LOG_FILE="${IMAGE_WORK_PATH}/log.txt"
TOOL_RES_PATH=${RUNNING_DIRECTORY}/res

NUM_JOBS=$(( `grep -c processor /proc/cpuinfo`+1 ))
EMERGE_PROXY="${FLAGS_proxy}"

if [ -z ${STAGE3_URL} ]; then
	STAGE3_URL="http://distribution.hexxeh.net/gentoo/stage3-i686-latest.tar.bz2"
fi

if [ -z ${PORTAGE_URL} ]; then
	PORTAGE_URL="http://distribution.hexxeh.net/gentoo/portage-latest.tar.bz2"
fi

if [ -z ${BINHOST_URL} ]; then
	BINHOST_URL="http://tinderbox.dev.gentoo.org/default-linux/x86"
fi

# Clean up old mounts
cleanup_mounts() {
	sleep 2
	umount -d -f ${IMAGE_WORK_PATH}/rootfs/dev/pts &>> ${LOG_FILE}
	umount -d -f ${IMAGE_WORK_PATH}/rootfs/dev/shm &>> ${LOG_FILE}
	umount -d -f ${IMAGE_WORK_PATH}/rootfs/dev &>> ${LOG_FILE}
	umount -d -f ${IMAGE_WORK_PATH}/rootfs/proc &>> ${LOG_FILE}
	umount -d -f -l  ${IMAGE_WORK_PATH}/rootfs/boot &>> ${LOG_FILE}
	umount -d -f -l ${IMAGE_WORK_PATH}/rootfs &>> ${LOG_FILE}
	sleep 2
}

# Handle some errors
handle_error() {
	echo "$1" 
	cd ${IMAGE_WORK_PATH}
	mv ${LOG_FILE} ${IMAGES_OUTPUT_PATH}/${BUILD_ID}.log || handle_error "Error moving log file"
	cleanup_mounts && rm -rf ${IMAGE_WORK_PATH}
	exit 1
}

echo "Step 1: Creating build working directory"
mkdir -p ${IMAGE_WORK_PATH} || handle_error "Error creating working directory"
cd ${IMAGE_WORK_PATH}
echo "" > ${LOG_FILE} || handle_error "Error creating log file"
cleanup_mounts

# Create disk iamge
BYTES_PER_CYLINDER=$(( 512*63*255 ))
CYLINDERS=$(( ${IMAGE_BYTES}/${BYTES_PER_CYLINDER} ))
REAL_IMAGE_BYTES=$(( (${CYLINDERS}+1)*${BYTES_PER_CYLINDER} ))
echo "Step 2: Creating image ${IMAGE_NAME}, size ${REAL_IMAGE_BYTES} bytes"
qemu-img create -f raw ${IMAGE_NAME} ${REAL_IMAGE_BYTES} &>> ${LOG_FILE} || handle_error "Error creating disk image file"

# Create partition table
echo "Step 3: Writing partition table"
echo -e "\x55\xaa" | dd bs=1 count=2 seek=510 of=${IMAGE_NAME} conv=notrunc &>> ${LOG_FILE}
LOOP_DEV_IMAGE=`losetup -f`
losetup --sizelimit ${REAL_IMAGE_BYTES} ${LOOP_DEV_IMAGE} ${IMAGE_NAME} &>> ${LOG_FILE} || handle_error "Error loop mounting disk image"

sfdisk ${LOOP_DEV_IMAGE} -H64 -S32 &>> ${LOG_FILE} << EOF
1,${BOOT_MEGABYTES},83,*
$(( 1+ ${BOOT_MEGABYTES} )),${SWAP_MEGABYTES},82,-
$(( 1+ ${BOOT_MEGABYTES} + ${SWAP_MEGABYTES} )),${ROOT_MEGABYTES},83,-
EOF
sleep 2
losetup -d ${LOOP_DEV_IMAGE} || handle_error "Error destroying disk image loop device"

# Mounting new root and boot
echo "Step 4: Creating filesystems"

LOOP_DEV_BOOT=`losetup -f`
losetup -o $(( 512 * 2048 )) --sizelimit $(( ${BOOT_MEGABYTES} * 1024 * 1024 )) ${LOOP_DEV_BOOT} ${IMAGE_NAME} || handle_error "Error loop mounting boot partition"
mkfs -t ext2 ${LOOP_DEV_BOOT} &>> ${LOG_FILE} || handle_error "Error formatting boot filesystem"
sleep 2
losetup -d ${LOOP_DEV_BOOT} || handle_error "Error destroying boot partition loop device"

LOOP_DEV_ROOT=`losetup -f`
losetup -o $(( ( 512 * 2048 )  + ( ${BOOT_MEGABYTES} * 1024 * 1024 ) + ( ${SWAP_MEGABYTES} * 1024 * 1024 ) )) --sizelimit $(( ${ROOT_MEGABYTES} * 1024 * 1024 )) ${LOOP_DEV_ROOT} ${IMAGE_NAME} || handle_error "Error loop mounting root partition"
mkfs -t ext3 ${LOOP_DEV_ROOT} &>> ${LOG_FILE} || handle_error "Error formatting root filesystem"
sleep 2
losetup -d ${LOOP_DEV_ROOT} || handle_error "Error destroying root partition loop device"

LOOP_DEV_SWAP=`losetup -f`
losetup -o $(( ( 512 * 2048 ) + ( ${BOOT_MEGABYTES} * 1024 * 1024 ) )) --sizelimit $(( ${SWAP_MEGABYTES} * 1024 * 1024 )) ${LOOP_DEV_SWAP} ${IMAGE_NAME} || handle_error "Error loop mounting swap partition"
mkswap ${LOOP_DEV_SWAP} &>> ${LOG_FILE}
sleep 2 
losetup -d ${LOOP_DEV_SWAP} || handle_error "Error destroying swap partition loop device"

echo "Step 5: Mounting fileystems"

mkdir -p rootfs
mount -o loop,offset=$(( ( 512 * 2048 )  + ( ${BOOT_MEGABYTES} * 1024 * 1024 ) + ( ${SWAP_MEGABYTES} * 1024 * 1024 ) )) ${IMAGE_NAME} rootfs  || handle_error "Error mounting root filesystem"
mkdir -p rootfs/boot
mount -o loop,offset=$(( 512 * 2048 )) ${IMAGE_NAME} rootfs/boot || handle_error "Error mounting boot filesystem"
cd rootfs

# Setup Gentoo

if [ ! -f ${TOOL_RES_PATH}/stage3.tar.bz2 ];
then
	wget ${STAGE3_URL} -O ${TOOL_RES_PATH}/stage3.tar.bz2 &>> ${LOG_FILE} || handle_error "Error downloading Stage3 tarball"
fi

echo "Step 6: Extracting Stage 3"
tar jxf ${TOOL_RES_PATH}/stage3.tar.bz2 &>> ${LOG_FILE} || handle_error "Error extracting Stage3 tarball"

if [ ! -f ${TOOL_RES_PATH}/portage-latest.tar.bz2 ];
then
	wget ${PORTAGE_URL} -O ${TOOL_RES_PATH}/portage-latest.tar.bz2 &>> ${LOG_FILE} || handle_error "Error downloading Portage snapshot"
fi

echo "Step 7: Extracting Portage snapshot"
cd usr
tar jxf ${TOOL_RES_PATH}/portage-latest.tar.bz2 &>> ${LOG_FILE} || handle_error "Error extracting Portage snapshot"

echo "Step 8: Setting up chroot"
cd ..
mount -t proc /proc proc || handle_error "Error mounting /proc"
mount --rbind /dev dev || handle_error "Error mounting /dev"
cp -L /etc/resolv.conf etc/resolv.conf || handle_error "Error setting up resolv.conf"

echo "Step 9: Setting up make.conf"
mkdir -p usr/portage/packages
echo "PORTAGE_BINHOST=\"${BINHOST_URL}\"" >> etc/make.conf
echo "PKGDIR=\"/usr/portage/packages\"" >> etc/make.conf
echo "FEATURES=\"${FEATURES}\"" >> etc/make.conf
echo "USE=\"${USE_FLAGS}\"" >> etc/make.conf
echo "MAKEOPTS=\"-j${NUM_JOBS}\"" >> etc/make.conf

if [[ ${OUTPUT_FORMAT} = "vbox" ]]; then
	echo 'INPUT_DEVICES="virtualbox evdev"' >> etc/make.conf
	echo 'VIDEO_CARDS="virtualbox"' >> etc/make.conf
fi

if [ -n $EMERGE_PROXY ]; then
	echo "Enabling HTTP proxy" &>> ${LOG_FILE}
	echo "http_proxy=\"${EMERGE_PROXY}\"" >> etc/make.conf
fi

echo "Step 10: Setting up package.use"
mkdir -p etc/portage
echo ${PACKAGE_USE} >> etc/portage/package.use

echo "Step 11: Setting up package.accept_keywords"
echo ${PACKAGE_ACCEPT_KEYWORDS} >> etc/portage/package.accept_keywords

echo "Step 12: Chroot setup"
linux32 chroot . env-update &>> ${LOG_FILE}
linux32 chroot . /bin/bash /etc/profile &>> ${LOG_FILE}

echo "Step 13: Setting timezone to ${TIMEZONE}"
linux32 chroot . cp /usr/share/zoneinfo/${TIMEZONE} /etc/localtime &>> ${LOG_FILE}

echo "Step 14: Setting keyboard layout to ${KEYBOARD}"
linux32 chroot . sed -i -e "s/keymap=.*/keymap=\"${KEYBOARD}\"/" /etc/conf.d/keymaps &>> ${LOG_FILE}

echo "Step 15: Setting hostname to ${HOSTNAME}"
linux32 chroot . /bin/bash -c "echo hostname='${HOSTNAME}' > /etc/conf.d/hostname" &>> ${LOG_FILE}
linux32 chroot . /bin/bash -c "echo 127.0.0.1 ${HOSTNAME}.local ${HOSTNAME} localhost > /etc/hosts" &>> ${LOG_FILE}

echo "Step 16: Copying new fstab"
cp ${TOOL_RES_PATH}/fstab etc/fstab &>> ${LOG_FILE}

echo "Step 17: Setting up networking"
echo 'config_eth0=( "dhcp" )' > etc/conf.d/net
cp etc/init.d/net.lo etc/init.d/net.eth0
linux32 chroot . rc-update add net.eth0 default &>> ${LOG_FILE}

echo "Step 18: Setting up kernel"
# If we got the flag, used a cached kernel to reduce build times for testing
if [[ ${FLAGS_cachedkernel} -eq ${FLAGS_TRUE} ]]; then
        echo "Using cached kernel" &>> ${LOG_FILE}
        cp ${TOOL_RES_PATH}/kernel boot/kernel || handle_error "Error copying cached kernel"
        mkdir lib/modules && cp  ${TOOL_RES_PATH}/modules.tar.gz lib/modules/ || handle_error "Error copying cached kernel modules"
        tar xvf lib/modules/modules.tar.gz -C lib/modules/ || handle_error "Error extracting cached kernel modules"
        KERNEL=`ls -l lib/modules/ | cut -d" " -f9 | tr "\n" " "`
        linux32 chroot . depmod -a ${KERNEL}
else
	echo "Downloading/installing kernel sources" &>> ${LOG_FILE}
	linux32 chroot . emerge gentoo-sources &>> ${LOG_FILE} || handle_error "Error emerging kernel sources"

	echo "Copying kernel configuration" &>> ${LOG_FILE}
	cp ${TOOL_RES_PATH}/kernelconfig usr/src/linux/.config || handle_error "Error copying kernel config"

	echo "Building kernel" &>> ${LOG_FILE}
	linux32 chroot . make -C /usr/src/linux -j${NUM_JOBS} &>> ${LOG_FILE} || handle_error "Error building kernel"

	echo "Installing kernel" &>> ${LOG_FILE}
	linux32 chroot . make -C /usr/src/linux modules_install &>> ${LOG_FILE} || handle_error "Error installing kernel modules"
	cp usr/src/linux/arch/i386/boot/bzImage boot/kernel &>> ${LOG_FILE} || handle_error "Error copying kernel"
fi

echo "Step 19: Setting root password"
linux32 chroot . /bin/bash -c "echo 'root:${ROOT_PASSWORD}' | chpasswd" &>> ${LOG_FILE} || handle_error "Error setting root password"

echo "Step 20: Processing packages list"
for PACKAGE in ${PACKAGES_LIST}; do
	echo "Installing ${PACKAGE}" &>> ${LOG_FILE}
	linux32 chroot . emerge --jobs=${NUM_JOBS} ${PACKAGE} &>> ${LOG_FILE} || handle_error "Error emerging ${PACKAGE}"
done

echo "Step 21: Adding default user"
linux32 chroot . useradd -g users -G lp,wheel,audio,cdrom,portage -m ${DEFAULT_USERNAME} || handle_error "Error adding default user"
linux32 chroot . /bin/bash -c "echo '${DEFAULT_USERNAME}:${DEFAULT_PASSWORD}' | chpasswd" &>> ${LOG_FILE} || handle_error "Error setting default user password"

if [[ ${OUTPUT_FORMAT} = "vbox" ]]
then
	echo "Installing VirtualBox additions/drivers"
	linux32 chroot . emerge xf86-video-virtualbox xf86-input-virtualbox virtualbox-guest-additions &>> ${LOG_FILE} || handle_error "Error emerging VirtualBox extras"
	linux32 chroot . rc-update add virtualbox-guest-additions default &>> ${LOG_FILE}
	mv etc/X11/xorg.conf etc/X11/xorg.conf.bak &>> ${LOG_FILE}
	linux32 chroot . usermod -a vboxguest ${DEFAULT_USERNAME}
fi

echo "Step 22: Cleaning up make.conf"
if [ -n ${EMERGE_PROXY} ]; then
	sed -i '/http_proxy/ d' etc/make.conf
fi
sed -i '/MAKEOPTS/ d' etc/make.conf

echo "Step 23: Installing extlinux"
extlinux --heads 255 --sectors 63 --install boot &>> ${LOG_FILE} || handle_error "Error installing extlinux"
dd if=/usr/lib/extlinux/mbr.bin of=../${IMAGE_NAME} conv=notrunc &>> ${LOG_FILE} || handle_error "Error copying extlinux MBR"
cp ${TOOL_RES_PATH}/extlinux.conf boot/ || handle_error "Error copying extlinux configuration"
cd  ..
cleanup_mounts
case "${OUTPUT_FORMAT}" in
	"raw" )
		echo "Already in raw format, not converting" &>> ${LOG_FILE}
		IMAGE_OUT="${BUILD_ID}.image"
		;;
	"vbox" )
		echo "Converting image from RAW to VDI" &>> ${LOG_FILE}
		qemu-img convert -O vdi ${IMAGE_NAME} ${BUILD_ID}.vdi &>> ${LOG_FILE} || handle_error "Error converting disk image to VDI format"
		rm -rf ${IMAGE_NAME}
		IMAGE_OUT="${BUILD_ID}.vdi"
		;;
	"vmware" )
		echo "Converting image from RAW to VMDK" &>> ${LOG_FILE}
		qemu-img convert -O vmdk ${IMAGE_NAME} ${BUILD_ID}.vmdk &>> ${LOG_FILE} || handle_error "Error converting disk image to VMDK format"
		rm -rf ${IMAGE_NAME}
		IMAGE_OUT="${BUILD_ID}.vmdk"
		;;
esac

mv ${IMAGE_OUT} ${IMAGES_OUTPUT_PATH}/${IMAGE_OUT} || handle_error "Error moving finished image"
mv ${LOG_FILE} ${IMAGES_OUTPUT_PATH}/${BUILD_ID}.log || handle_error "Error moving log file"
LOG_FILE="${IMAGES_OUTPUT_PATH}/${BUILD_ID}.log"
rm -rf ${IMAGE_WORK_PATH} || handle_error "Error removing working directory"

if [[ ${FLAGS_compress} -eq ${FLAGS_TRUE} ]]; then
	cd ${IMAGES_OUTPUT_PATH}
	tar czvf "${BUILD_ID}.tar.gz" "${IMAGE_OUT}" &>> ${LOG_FILE} || handle_error "Error compressing image"
	IMAGE_OUT="${BUILD_ID}.tar.gz"
fi

echo "Step 24: Image build completed!"
echo "Your image is here: ${IMAGES_OUTPUT_PATH}/${IMAGE_OUT}"
echo "Your log file is here: ${IMAGES_OUTPUT_PATH}/${BUILD_ID}.log"