#!/bin/bash
# This script downloads the latest Linux kernel and rootfs images from Firecracker CI, patches the rootfs with an SSH key, and creates an ext4 filesystem image.

required_commands=(
    uname
    curl
    basename
    grep
    sort
    tail
    wget
    unsquashfs
    ssh-keygen
    sudo
    chown
    truncate
    mkfs.ext4
    e2fsck
)

missing_commands=()
for cmd in "${required_commands[@]}"; do
    if ! command -v "$cmd" >/dev/null 2>&1; then
        missing_commands+=("$cmd")
    fi
done

if [ ${#missing_commands[@]} -ne 0 ]; then
    echo "ERROR: Missing required program(s): ${missing_commands[*]}" >&2
    echo "Install missing dependencies and rerun." >&2
    exit 1
fi

ARCH="$(uname -m)"
release_url="https://github.com/firecracker-microvm/firecracker/releases"
latest_version=$(basename "$(curl -fsSLI -o /dev/null -w "%{url_effective}" "${release_url}/latest")")
CI_VERSION=${latest_version%.*}
latest_kernel_key=$(curl "http://spec.ccfc.min.s3.amazonaws.com/?prefix=firecracker-ci/$CI_VERSION/$ARCH/vmlinux-&list-type=2" \
    | grep -oP "(?<=<Key>)(firecracker-ci/$CI_VERSION/$ARCH/vmlinux-[0-9]+\.[0-9]+\.[0-9]{1,3})(?=</Key>)" \
    | sort -V | tail -1)

# Download a linux kernel binary
wget "https://s3.amazonaws.com/spec.ccfc.min/${latest_kernel_key}"

latest_ubuntu_key=$(curl "http://spec.ccfc.min.s3.amazonaws.com/?prefix=firecracker-ci/$CI_VERSION/$ARCH/ubuntu-&list-type=2" \
    | grep -oP "(?<=<Key>)(firecracker-ci/$CI_VERSION/$ARCH/ubuntu-[0-9]+\.[0-9]+\.squashfs)(?=</Key>)" \
    | sort -V | tail -1)
ubuntu_version=$(basename "$latest_ubuntu_key" .squashfs | grep -oE '[0-9]+\.[0-9]+')

# Download a rootfs from Firecracker CI
wget -O "ubuntu-${ubuntu_version}.squashfs.upstream" "https://s3.amazonaws.com/spec.ccfc.min/$latest_ubuntu_key"

# The rootfs in our CI doesn't contain SSH keys to connect to the VM
# For the purpose of this demo, let's create one and patch it in the rootfs
unsquashfs "ubuntu-${ubuntu_version}.squashfs.upstream"
ssh-keygen -f id_rsa -N ""
cp -v id_rsa.pub squashfs-root/root/.ssh/authorized_keys
mv -v id_rsa "./ubuntu-${ubuntu_version}.id_rsa"
# create ext4 filesystem image
sudo chown -R root:root squashfs-root
truncate -s 1G "ubuntu-${ubuntu_version}.ext4"
sudo mkfs.ext4 -d squashfs-root -F "ubuntu-${ubuntu_version}.ext4"

# Verify everything was correctly set up and print versions
echo
echo "The following files were downloaded and set up:"
kernel_files=(./vmlinux-*)
if [ -e "${kernel_files[0]}" ]; then
    last_kernel_index=$((${#kernel_files[@]} - 1))
    KERNEL="${kernel_files[$last_kernel_index]}"
    echo "Kernel: $KERNEL"
else
    echo "ERROR: Kernel image does not exist"
fi

rootfs_files=(./*.ext4)
if [ -e "${rootfs_files[0]}" ]; then
    last_rootfs_index=$((${#rootfs_files[@]} - 1))
    ROOTFS="${rootfs_files[$last_rootfs_index]}"
    e2fsck -fn "$ROOTFS" &>/dev/null && echo "Rootfs: $ROOTFS" || echo "ERROR: $ROOTFS is not a valid ext4 fs"
else
    echo "ERROR: Rootfs image does not exist"
fi

key_files=(./*.id_rsa)
if [ -e "${key_files[0]}" ]; then
    last_key_index=$((${#key_files[@]} - 1))
    KEY_NAME="${key_files[$last_key_index]}"
    echo "SSH Key: $KEY_NAME"
else
    echo "ERROR: SSH key does not exist"
fi