#! /bin/bash

# fai-diskimage - create a disk image for booting a VM

# This script is part of FAI (Fully Automatic Installation)
#
# Copyright (C) 2016-2025 Thomas Lange, lange@cs.uni-koeln.de
# Universitaet zu Koeln

die() {

    local e=$1   # first parameter is the exit code
    shift

    echo "ERROR: $*" >&2   # print error message
    exit $e
}

cleanup() {

    local dir
    for dir in $(mount | awk '{print $3}'| grep $mnt | sort -r); do
	# try umount several times, because some postinst jobs may still run and use the filesystems
	for i in {1..15}; do
            umount $dir >&/dev/null && break
            [ $((i % 3)) -eq 0 ] && echo "Waiting for background jobs to finish."
            [ $i -gt 12 ] && echo "Still cannot umount $dir. Will try again."
	    sleep $i
	done
	echo "ERROR: fai-diskimage cannot umount $dir" >&2
    done

    # call zerofree for ext2/3/4 devices if available
    if [ -f /var/run/fai/zerofree.$$ ] && [ -x /usr/sbin/zerofree ]; then
        . /var/run/fai/zerofree.$$ 2>/dev/null
        rm /var/run/fai/zerofree.$$
    fi
    # if FAI created a volume group, we can remove it after the loop device is removed
    if [ -f /var/run/fai/vgremove.$$ ]; then
	. /var/run/fai/vgremove.$$ 2>/dev/null
	rm /var/run/fai/vgremove.$$
    fi
    rm -rf $mnt
    losetup -d $loop
    if [ -f /var/run/fai/FAI_INSTALLATION_IN_PROGRESS ]; then
	if pgrep -F /var/run/fai/FAI_INSTALLATION_IN_PROGRESS; then
	    :
	else
	    rm /var/run/fai/FAI_INSTALLATION_IN_PROGRESS
	fi
    fi
    if [ -d /sys/modules/loop ]; then
        rmmod loop
    fi
}

usage() {

    echo "Usage: $0 name.raw

Create a disk image name.raw using FAI and a list of FAI classes.
This can be used for a virtual machine or a cloud instance. If you
use another suffix the image will be converted. Following formats are
supported: .raw.xz, .raw.zst, .qcow2, .vdi, .vhdx, .vmdk, .simg.

   -h|--help      	    print help
   -v|--verbose   	    be verbose
   -N|--new                 execute scripts class/[0-9]* for defining classes
   -c|--class <class,...>   define list of FAI classes
   -C|--cfdir <dir>         Use dir for reading the config files (default /etc/fai)
   -S|--size  <size>        size of raw image (suffixes k M G T are supported)
   -s|--cspace  <uri>       location of the config space
   -u|--hostname <name>     set hostname to name
"
    exit $1
}

check_commands() {

    local error=0
    if ! command -v qemu-img >&/dev/null; then
	echo "qemu-img not found. Install the package qemu-utils."
	error=1
    fi
    if ! command -v setup-storage >&/dev/null; then
	echo "setup-storage not found. Install the package fai-setup-storage."
	error=1
    fi
    if [ $error -eq 1 ]; then
	die 5 "Aborted."
    fi
}

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

call="$0 $*"
TEMP=$(getopt -o C:NS:s:u:hvc: --long new,cfdir:,cspace:,hostname:,class:,size:,help,verbose -n "$0" -- "$@")
if [ $? != 0 ] ; then die 6 "Wrong option. Terminating." >&2 ; fi
# Note the quotes around `$TEMP': they are essential!
eval set -- "$TEMP"
unset TEMP

verbose=0
convert=1

while true ; do
    case "$1" in
        -h|--help)
            usage 0 ;;
        -v|--verbose)
            export verbose=1
            shift ;;
        -c|--class)
            export classes="-c $2"
            shift 2 ;;
        -C|--cfdir)
            cfdir="-C $2 "
            shift 2 ;;
        -N|--new)
            renew="-N"
            shift ;;
        -S|--size)
            size=$2
            shift 2 ;;
        -s|--cspace)
            cspace=$2
            shift 2 ;;
        -u|--hostname)
            export hname=$2
            shift 2 ;;
        --)
            shift
            break ;;
         *)
            die 1 "$0: command line parsing error ! $*" >&2 ;;
    esac
done

# check options, set defaults

[ "$1" ] || usage
image=$1

iname=${image%.*}   # strip last suffix
iname2=${iname%.*}  # strip next to last suffix
rawname="$iname.raw"

case "$image" in
    *.raw.xz)
        convert=2
	rawname=$iname
	export XZ_OPT=${XZ_OPT-"-1 -T0"}
	cmd="xz $rawname"
        ;;

    *.raw.zst)
        convert=2
	rawname=$iname
	cmd="zstd --rm -q -9 -T0 $rawname"
        ;;

    *.raw)
        convert=0
        ;;

    *.qcow2)
        copt="-O qcow2 -c -o compression_type=zstd "
	qcowname="$iname.qcow2"
	;;

    *.simg) convert=2
	  cmd="img2simg $rawname $iname.simg"
	  if ! command -v img2simg >&/dev/null; then
	      echo "img2simg not found. Install the package img2simg."
	      error=1
	  fi
	  ;;
    *.vdi) copt="-O vdi"
	 qcowname="$iname.vdi"
	   ;;
    *.vmdk) copt="-O vmdk"
	  qcowname="$iname.vmdk"
	  ;;
    *.vhdx) copt="-O vhdx"
	  qcowname="$iname.vhdx"
	  ;;
    *) die 8 "Unknown suffix. Please use raw, raw.zst, raw.xz, qcow2, vdi, vmdk or vhdx, simg."
esac

if [ -z "$classes" ] && [ -z "$renew" ]; then
    die 7 "No classes are defined. Use -c or -N."
fi

: ${cspace:=/srv/fai/config}
# if cspace starts with /, add prefix file://
uri=${cspace/#\//file:///}

: ${size:=800M}
: ${hname:=debian.example.com}

# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -

check_commands

# check root
if [ $(id -u) != "0" ]; then
    die 1 "Run this program as root."
fi

if command -v vgs >&/dev/null; then
    export SS_IGNORE_VG=$(vgs --unbuffered --noheadings -o name --rows)
fi

# create empty disk image, loop setup, temp mount point
rm -f $rawname
qemu-img create $rawname $size
loop=$(losetup -P -f --show $rawname)
loopdev=${loop/\/dev\//}
export disklist=$loopdev
mnt=$(mktemp -d -t fai-diskimage.XXXXXX)

trap "cleanup" EXIT

LC_ALL=C fai ${cfdir}${renew} -u $hname -s $uri $classes install $mnt
error=$?
cleanup

trap - EXIT

# convert if needed
if [ $convert -eq 1 ] || [ $convert -eq 3 ]; then
    [ $verbose -eq 1 ] && echo "Converting $rawname to $qcowname"
    qemu-img convert -f raw $rawname $copt $qcowname
    rm $rawname
fi
if [ $convert -ge 2 ]; then
    [ $verbose -eq 1 ] && echo "Converting to $image"
    $cmd
fi
echo -n "Size of disk image and filename: "; du -h $image
echo "Image created by: $call"
exit $error