# vi: ts=4 expandtab syntax=sh

imagesize=${IMAGE_SIZE:-$((2252*1024**2))}  # 2.2G (the current size we ship)
fs_label="${FS_LABEL:-rootfs}"

rootfs_dev_mapper=
loop_device=
loop_raw=
backing_img=

clean_loops() {
    local kpartx_ret
    local kpartx_stdout

    if [ -n "${backing_img}" ]; then
        # sync before removing loop to avoid "Device or resource busy" errors
        sync
        kpartx_ret=""
        kpartx_stdout=$(kpartx -v -d "${backing_img}") || kpartx_ret=$?
        echo "$kpartx_stdout"
        if [ -n "$kpartx_ret" ]; then
            if echo "$kpartx_stdout" | grep -q "loop deleted"; then
                echo "Suppressing kpartx returning error (#860894)"
            else
                exit $kpartx_ret
            fi
        fi
        unset backing_img
    fi

    if [ -z "${rootfs_dev_mapper}" ]; then
        return 0
    fi

    unset loop_device
    unset loop_raw
    unset rootfs_dev_mapper
}

create_empty_disk_image() {
    # Prepare an empty disk image
    dd if=/dev/zero of="$1" bs=1 count=0 seek="${imagesize}"
}

make_ext4_partition() {
    device="$1"
    label=${fs_label:+-L "${fs_label}"}
    mkfs.ext4 -F -b 4096 -i 8192 -m 0 ${label} -E resize=536870912 "$device"
}

mount_image() {
    trap clean_loops EXIT
    backing_img="$1"
    local rootpart="$2"
    kpartx_mapping="$(kpartx -s -v -a ${backing_img})"

    # Find the loop device
    loop_p1="$(echo -e ${kpartx_mapping} | head -n1 | awk '{print$3}')"
    loop_device="/dev/loop$(echo ${loop_p1} | cut -b5)"
    if [ ! -b ${loop_device} ]; then
        echo "unable to find loop device for ${backing_img}"
        exit 1
    fi

    # Find the rootfs location
    rootfs_dev_mapper="/dev/mapper/${loop_p1%%[0-9]}${rootpart}"
    if [ ! -b "${rootfs_dev_mapper}" ]; then
        echo "${rootfs_dev_mapper} is not a block device";
        exit 1
    fi

    # Add some information to the debug logs
    echo "Mounted disk image ${backing_img} to ${rootfs_dev_mapper}"
    blkid ${rootfs_dev_mapper}

    return 0
}

setup_mountpoint() {
    local mountpoint="$1"

    mount --rbind --make-rslave /dev "$mountpoint/dev"
    mount proc-live -t proc "$mountpoint/proc"
    mount sysfs-live -t sysfs "$mountpoint/sys"
    mount -t tmpfs none "$mountpoint/tmp"
    mount -t tmpfs none "$mountpoint/var/lib/apt"
    mount -t tmpfs none "$mountpoint/var/cache/apt"
    mv "$mountpoint/etc/resolv.conf" resolv.conf.tmp
    cp /etc/resolv.conf "$mountpoint/etc/resolv.conf"
    chroot "$mountpoint" apt-get update

}

teardown_mountpoint() {
    # Reverse the operations from setup_mountpoint
    local mountpoint="$1"

    # ensure we have exactly one trailing slash, and escape all slashes for awk
    mountpoint_match=$(echo "$mountpoint" | sed -e's,/$,,; s,/,\\/,g;')'\/'
    # sort -r ensures that deeper mountpoints are unmounted first
    for submount in $(awk </proc/self/mounts "\$2 ~ /$mountpoint_match/ \
                      { print \$2 }" | LC_ALL=C sort -r); do
        umount $submount
    done
    mv resolv.conf.tmp "$mountpoint/etc/resolv.conf"
}

mount_partition() {
    partition="$1"
    mountpoint="$2"

    mount "$partition" "$mountpoint"
    setup_mountpoint "$mountpoint"
}

mount_overlay() {
    lower="$1"
    upper="$2"
    work="$2/../work"
    path="$3"

    mkdir -p "$work"
    mount -t overlay overlay \
	-olowerdir="$lower",upperdir="$upper",workdir="$work" \
	"$path"
}

mount_disk_image() {
    local disk_image=${1}
    local mountpoint=${2}
    mount_image ${disk_image} 1
    mount_partition "${rootfs_dev_mapper}" $mountpoint

    local uefi_dev="/dev/mapper${loop_device///dev/}p15"
    if [ -b ${uefi_dev} -a -e $mountpoint/boot/efi ]; then
        mount "${uefi_dev}" $mountpoint/boot/efi
    fi

    # This is needed to allow for certain operations
    # such as updating grub and installing software
    cat > $mountpoint/usr/sbin/policy-rc.d << EOF
#!/bin/sh
# ${IMAGE_STR}
echo "All runlevel operations denied by policy" >&2
exit 101
EOF
    chmod 0755 $mountpoint/usr/sbin/policy-rc.d

}

umount_partition() {
    local mountpoint=${1}
    teardown_mountpoint $mountpoint
    umount -R $mountpoint
    udevadm settle

    if [ -n "${rootfs_dev_mapper}" -a -b "${rootfs_dev_mapper}" ]; then
        # buildd's don't have /etc/mtab symlinked
        # /etc/mtab is needed in order zerofree space for ext4 filesystems
        [ -e /etc/mtab ] || ln -s /proc/mounts /etc/mtab

        # both of these are likely overkill, but it does result in slightly
        # smaller ext4 filesystem
        e2fsck -y -E discard ${rootfs_dev_mapper}
        zerofree ${rootfs_dev_mapper}
    fi
}

umount_disk_image() {
    mountpoint="$1"

    local uefi_dev="/dev/mapper${loop_device///dev/}p15"
    if [ -e "$mountpoint/boot/efi" -a -b "$uefi_dev" ]; then
        umount --detach-loop "$mountpoint/boot/efi"
    fi

    if [ -e $mountpoint/usr/sbin/policy-rc.d ]; then
        rm $mountpoint/usr/sbin/policy-rc.d
    fi
    umount_partition $mountpoint
    clean_loops
}

modify_vmdk_header() {
    # Modify the VMDK headers so that both VirtualBox _and_ VMware can
    # read the vmdk and import them.

    vmdk_name="${1}"
    descriptor=$(mktemp)
    newdescriptor=$(mktemp)

    # Extract the vmdk header for manipulation
    dd if="${vmdk_name}" of="${descriptor}" bs=1 skip=512 count=1024

    # The sed lines below is where the magic is. Specifically:
    #   ddb.toolsVersion: sets the open-vm-tools so that VMware shows
    #       the tooling as current
    #   ddb.virtualHWVersion: set the version to 7, which covers most
    #       current versions of VMware
    #   createType: make sure its set to stream Optimized
    #   remove the vmdk-stream-converter comment and replace with
    #       # Disk DescriptorFile. This is needed for Virtualbox
    #   remove the comments from vmdk-stream-converter which causes
    #       VirtualBox and others to fail VMDK validation

    sed -e 's|# Description file.*|# Disk DescriptorFile|' \
        -e '/# Believe this is random*/d' \
        -e '/# Indicates no parent/d' \
        -e '/# The Disk Data Base/d' \
        -e 's|ddb.comment.*|ddb.toolsVersion = "2147483647"|' \
            "${descriptor}" > "${newdescriptor}"

    # The header is cannot be bigger than 1024
    expr $(stat --format=%s ${newdescriptor}) \< 1024 > /dev/null 2>&1 || {
        echo "descriptor is too large, VMDK will be invalid!"; exit 1; }

    # Overwrite the vmdk header with our new, modified one
    dd conv=notrunc,nocreat \
        if="${newdescriptor}" of="${vmdk_name}" \
        bs=1 seek=512 count=1024

    rm ${descriptor} ${newdescriptor}
}

create_vmdk() {
    # There is no real good way to create a _compressed_ VMDK using open source
    # tooling that works across multiple VMDK-capable platforms. This functions
    # uses vmdk-stream-converter and then calls modify_vmdk_header to produce a
    # compatible VMDK.

    src="$1"
    destination="$2"
    size="${3:-10240}"

    streamconverter="VMDKstream"
    scratch_d=$(mktemp -d)
    cp ${src} ${scratch_d}/resize.img

    truncate --size=${size}M ${scratch_d}/resize.img
    python -m ${streamconverter} ${scratch_d}/resize.img ${destination}
    modify_vmdk_header ${destination}

    qemu-img info ${destination}
    rm -rf ${scratch_d}
}

create_derivative() {
    # arg1 is the disk type
    # arg2 is the new name
    unset derivative_img
    case ${1} in
           uefi) disk_image="binary/boot/disk-uefi.ext4";
                 dname="${disk_image//-uefi/-$2-uefi}";;
              *) disk_image="binary/boot/disk.ext4";
                 dname="${disk_image//.ext4/-$2.ext4}";;
    esac

    if [ ! -e ${disk_image} ]; then
        echo "Did not find ${disk_image}!"; exit 1;
    fi

    cp ${disk_image} ${dname}
    export derivative_img=${dname}
}

convert_to_qcow2() {
    src="$1"
    destination="$2"
    qemu-img convert -c -O qcow2 -o compat=0.10 "$src" "$destination"
    qemu-img info "$destination"
}

replace_grub_root_with_label() {
    # When update-grub is run, it will detect the disks in the build system.
    # Instead, we want grub to use the right labelled disk
    CHROOT_ROOT="$1"

    # If boot by partuuid has been requested, don't override.
    if [ -f $CHROOT_ROOT/etc/default/grub.d/40-force-partuuid.cfg ] && \
           grep -q ^GRUB_FORCE_PARTUUID= $CHROOT_ROOT/etc/default/grub.d/40-force-partuuid.cfg
    then
        return 0
    fi
    sed -i -e "s,root=[^ ]*,root=LABEL=${fs_label}," \
        "$CHROOT_ROOT/boot/grub/grub.cfg"
}


# When running update-grub in a chroot on a build host, we don't want it to
# probe for disks or probe for other installed OSes.  Extract common
# diversion wrappers, so this isn't reinvented differently for each image.
divert_grub() {
	CHROOT_ROOT="$1"

	chroot "$CHROOT_ROOT" dpkg-divert --local \
		--rename /usr/sbin/grub-probe
	chroot "$CHROOT_ROOT" touch /usr/sbin/grub-probe
	chroot "$CHROOT_ROOT" chmod +x /usr/sbin/grub-probe

	chroot "$CHROOT_ROOT" dpkg-divert --local \
		--divert /etc/grub.d/30_os-prober.dpkg-divert \
		--rename /etc/grub.d/30_os-prober

	# Divert systemd-detect-virt; /etc/kernel/postinst.d/zz-update-grub
	# no-ops if we are in a container, and the launchpad farm runs builds
	# in lxd.  We therefore pretend that we're never in a container (by
	# exiting 1).
	chroot "$CHROOT_ROOT" dpkg-divert --local \
		--rename /usr/bin/systemd-detect-virt
	echo "exit 1" > "$CHROOT_ROOT"/usr/bin/systemd-detect-virt
	chmod +x "$CHROOT_ROOT"/usr/bin/systemd-detect-virt
}

undivert_grub() {
	CHROOT_ROOT="$1"

	chroot "$CHROOT_ROOT" rm /usr/sbin/grub-probe
	chroot "$CHROOT_ROOT" dpkg-divert --remove --local \
		--rename /usr/sbin/grub-probe

	chroot "$CHROOT_ROOT" dpkg-divert --remove --local \
		--divert /etc/grub.d/30_os-prober.dpkg-divert \
		--rename /etc/grub.d/30_os-prober

	rm "$CHROOT_ROOT"/usr/bin/systemd-detect-virt
	chroot "$CHROOT_ROOT" dpkg-divert --remove --local \
		--rename /usr/bin/systemd-detect-virt
}