# vi: ts=4 expandtab syntax=sh

#imagesize=${IMAGE_SIZE:-$((2252*1024**2))}  # 2.2G (the current size we ship)
imagesize=${IMAGE_SIZE:-2361393152}  # 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_p1%p[0-9]*}"
    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_resolvconf() {
    local mountpoint=${1}
    mv "${mountpoint}/etc/resolv.conf" resolv.conf.tmp
    cp /etc/resolv.conf "${mountpoint}/etc/resolv.conf"
}

recover_resolvconf() {
    local mountpoint=${1}
    mv resolv.conf.tmp "${mountpoint}/etc/resolv.conf"
}

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"
    setup_resolvconf "${mountpoint}"
    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
    recover_resolvconf "${mountpoint}"
}

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

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


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
        # zero fill free space in UEFI partition
        cat < /dev/zero > "$mountpoint/boot/efi/bloat_file" 2> /dev/null || true
        rm "$mountpoint/boot/efi/bloat_file"
        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. The vodoo here is _not_ documented
    # anywhere....so this will have to do. This is undocumented vodoo
    # that has been learned by the Cloud Image team.

    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="/usr/share/pyshared/VMDKstream.py"
    scratch_d=$(mktemp -d)
    cp ${src} ${scratch_d}/resize.img

    truncate --size=${size}M ${scratch_d}/resize.img
    python ${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
}

release_ver() {
    # Return the release version number
    distro-info --series="$LB_DISTRIBUTION" -r | awk '{ print $1 }'
}

_snap_preseed() {
    # Download the snap/assertion and add to the preseed
    local CHROOT_ROOT=$1
    local SNAP=$2
    local SNAP_NAME=${SNAP%/*}
    local CHANNEL=${3:?Snap channel must be specified}

    local seed_dir="$CHROOT_ROOT/var/lib/snapd/seed"
    local snaps_dir="$seed_dir/snaps"
    local seed_yaml="$seed_dir/seed.yaml"
    local assertions_dir="$seed_dir/assertions"

    # Download the snap & assertion
    local snap_download_failed=0
    chroot $CHROOT_ROOT sh -c "
        set -x;
        cd /var/lib/snapd/seed;
        SNAPPY_STORE_NO_CDN=1 snap download \
            --channel=$CHANNEL \"$SNAP_NAME\"" || snap_download_failed=1
    if [ $snap_download_failed = 1 ] ; then
        echo "If the channel ($CHANNEL) includes '*/ubuntu-##.##' track per "
        echo "Ubuntu policy (ex. stable/ubuntu-18.04) the publisher will need "
        echo "to temporarily create the channel/track to allow fallback during"
        echo "download (ex. stable/ubuntu-18.04 falls back to stable if the"
        echo "prior had been created in the past)."
        exit 1
    fi

    mv -v $seed_dir/*.assert $assertions_dir
    mv -v $seed_dir/*.snap $snaps_dir

    # Add the snap to the seed.yaml
    ! [ -e $seed_yaml ] && echo "snaps:" > $seed_yaml
    cat <<EOF >> $seed_yaml
  -
    name: ${SNAP_NAME}
    channel: ${CHANNEL}
EOF

    case ${SNAP} in */classic) echo "    classic: true" >> $seed_yaml;; esac

    echo -n "    file: " >> $seed_yaml
    (cd $snaps_dir; ls -1 ${SNAP_NAME}_*.snap) >> $seed_yaml
}

snap_prepare_assertions() {
    # Configure basic snapd assertions
    local CHROOT_ROOT=$1
    # A colon-separated string of brand:model to be used for the image's model
    # assertion
    local CUSTOM_BRAND_MODEL=$2

    local seed_dir="$CHROOT_ROOT/var/lib/snapd/seed"
    local snaps_dir="$seed_dir/snaps"
    local assertions_dir="$seed_dir/assertions"
    local model_assertion="$assertions_dir/model"
    local account_key_assertion="$assertions_dir/account-key"
    local account_assertion="$assertions_dir/account"

    mkdir -p "$assertions_dir"
    mkdir -p "$snaps_dir"

    local brand="$(echo $CUSTOM_BRAND_MODEL | cut -d: -f 1)"
    local model="$(echo $CUSTOM_BRAND_MODEL | cut -d: -f 2)"

    if ! [ -e "$model_assertion" ] ; then
        snap known --remote model series=16 \
            model=$model brand-id=$brand \
            > "$model_assertion"
    fi

    if ! [ -e "$account_key_assertion" ] ; then
        local account_key=$(sed -n -e's/sign-key-sha3-384: //p' \
            < "$model_assertion")
        snap known --remote account-key \
            public-key-sha3-384="$account_key" \
            > "$account_key_assertion"
    fi


    if ! [ -e "$account_assertion" ] ; then
        local account=$(sed -n -e's/account-id: //p' < "$account_key_assertion")
        snap known --remote account account-id=$account \
            > "$account_assertion"
    fi
}

snap_prepare() {
    # Configure basic snapd assertions and pre-seeds the 'core' snap
    local CHROOT_ROOT=$1
    # Optional. If set, should be a colon-separated string of brand:model to be
    # used for the image's model assertion
    local CUSTOM_BRAND_MODEL=${2:-generic:generic-classic}

    local seed_dir="$CHROOT_ROOT/var/lib/snapd/seed"
    local snaps_dir="$seed_dir/snaps"

    snap_prepare_assertions "$CHROOT_ROOT" "$CUSTOM_BRAND_MODEL"

    # Download the core snap
    if ! [ -f $snaps_dir/core_[0-9]*.snap ] ; then
        _snap_preseed $CHROOT_ROOT core stable
    fi
}

snap_preseed() {
    # Preseed a snap in the image
    local CHROOT_ROOT=$1
    local SNAP=$2
    # Per Ubuntu policy, all seeded snaps (with the exception of the core
    # snap) must pull from stable/ubuntu-$(release_ver) as their channel.
    local CHANNEL=${3:-"stable/ubuntu-$(release_ver)"}

    snap_prepare $CHROOT_ROOT
    _snap_preseed $CHROOT_ROOT $SNAP $CHANNEL
}