#!/bin/bash

set -e

usage() {
    cat 1>&2 <<EOF
Please run with either "debian" or "ubuntu" as argument.
EOF
}

if [ "$UID" -ne 0 ]; then
    echo "We'll need root for this..." 1>&2
    exec sudo "$0" "$@"
fi

while [ -n "$1" ]; do
    case "$1" in
        "ubuntu")
            GENERATE_VARIANT=generate_ubuntu
            ;;
        "debian")
            GENERATE_VARIANT=generate_debian
            ;;
        *)
            usage
            exit 1
            ;;
    esac
    shift
done

cleanup_losetup() {
    set +e
    for dev in ${LO_DEVICE}p*; do
        umount $dev
    done
    losetup -d $LO_DEVICE
    rmdir tmp-p1
    rmdir tmp-p2
}

generate_debian() {
    local -r image="Debian-11-x86-64.img"

    mkosi --root-size=2G --distribution=debian --release=bullseye --format=gpt_ext4 --bootable --checksum \
          --password password --package=openssh-server,dhcpcd5 --package grub-efi-amd64-signed \
          --package shim-signed --package lsb-release --output="$image" build

    post_process_image "$image"

    echo "Image successfully generated!" 1>&2
}

generate_ubuntu() {
    local -r image="Ubuntu-Focal-x86-64.img"

    mkosi --root-size=2G --distribution=ubuntu --release=focal --format=gpt_ext4 --bootable --checksum \
          --password password --package=openssh-server,dhcpcd5 --package grub-efi-amd64-signed \
          --package shim-signed --package lsb-release --output="$image" build

    post_process_image "$image"

    echo "Image successfully generated!" 1>&2
}

post_process_image() {
    local -r image="$1"

    mkdir -p tmp-p1
    mkdir -p tmp-p2

    LO_DEVICE=$(losetup --find --show --partscan "$image")
    trap cleanup_losetup EXIT
    mount ${LO_DEVICE}p1 tmp-p1
    mount ${LO_DEVICE}p2 tmp-p2

    pre_tweaks tmp-p1 tmp-p2

    create_grub_regeneration_service tmp-p1 tmp-p2
    umount tmp-p1
    umount tmp-p2
    regenerate_grub_live "$image"
    mount ${LO_DEVICE}p1 tmp-p1
    mount ${LO_DEVICE}p2 tmp-p2

    post_tweaks tmp-p1 tmp-p2
}

pre_tweaks() {
    local -r boot="$1"
    local -r root="$2"

    # Fstab is missing for some reason. I'm not exactly sure why systemd-boot
    # works without this, and GRUB doesn't.
    cat > "$root/etc/fstab" <<EOF
/dev/root / auto defaults 0 0
/dev/sda1 /boot/efi auto defaults 0 0
EOF

    # Real installers create a /etc/default/grub file with a distributor in
    # them.
    if [ ! -e "$root/etc/default/grub" ]; then
        mkdir -p $root/etc/default
        echo 'GRUB_DISTRIBUTOR=`lsb_release -i -s 2> /dev/null`' > $root/etc/default/grub
    fi

    sed -E -i -e 's/^#? *PermitRootLogin .*/PermitRootLogin yes/' $root/etc/ssh/sshd_config
}

post_tweaks() {
    local -r boot="$1"
    local -r root="$2"

    # Delete systemd-boot, which isn't normally present in images that were
    # installed with OS installers, at least not at the time of writing.
    rm -rf "$boot/EFI/systemd"

    # Also replace bootx64.efi, which is the default bootloader. Mkosi installs
    # systemd-bootx86.efi, but we want the shim.
    rm -f "$boot/EFI/BOOT/*"
    mkdir -p "$boot/EFI/BOOT"
    cp "$root/usr/lib/shim/shimx64.efi.signed" "$boot/EFI/BOOT/BOOTX64.EFI"
}

# Unfortunately installing grub scripts is something which is not really
# possible when offline. This is something which is easier with systemd-boot, so
# longterm GRUB will probably follow, or systemd-boot will take over. Anyway,
# let's do it by using a systemd service to perform the job, and then shut down.
create_grub_regeneration_service() {
    local -r boot="$1"
    local -r root="$2"

    cat > "$root/etc/systemd/system/mender-regenerate-grub-and-shutdown.service" <<EOF
[Unit]
Description=Regenerate grub scripts, disable itself and then shut down.

[Service]
Type=oneshot
ExecStart=sh -c "grub-install && update-grub && systemctl disable mender-regenerate-grub-and-shutdown.service && poweroff"
EOF

    ln -sf "$root/etc/systemd/system/mender-regenerate-grub-and-shutdown.service" "$root/etc/systemd/system/multi-user.target.wants/"
}

regenerate_grub_live() {
    local -r image="$1"

    local -r nvvars=$(mktemp)
    dd if=/dev/zero of="$nvvars" bs=1M count=1

    local ret=0
    for maybe_kvm in -enable-kvm ""; do
        ret=0
        echo "Generating GRUB boot files live..." 1>&2
        qemu-system-x86_64 \
            $maybe_kvm \
            -drive file="$image",if=ide,format=raw \
            -drive file=/usr/share/OVMF/OVMF_CODE.fd,if=pflash,format=raw,unit=0,readonly=on \
            -drive file="$nvvars",if=pflash,format=raw,unit=1 \
            -display vnc=:23 \
            -m 512 \
            || ret=$?
        [ $ret -eq 0 ] && break
    done

    return $ret
}

$GENERATE_VARIANT