#!/bin/bash

set -e

. /usr/share/debconf/confmodule

clean_up ()
{
    cd /
    if [ "$CLEANUP" -gt 2 ]; then
        rm -fr "$TMPDIR"/upper "$TMPDIR"/initrd "$TMPDIR"/initrd.lz "$TMPDIR"/microcode.cpio
    fi
    if [ "$CLEANUP" -gt 1 ]; then
        umount "$TMPDIR"/chroot/dev "$TMPDIR"/chroot/sys "$TMPDIR"/chroot/proc "$TMPDIR"/chroot "$TMPDIR"/lower
        rm -fr "$TMPDIR"
    fi
    if [ "$CLEANUP" -gt 0 ] && [ -n "$DISK" ] && [ -d "$DISK" ]; then
        umount "$DISK"
        rmdir "$DISK"
    fi
}

generate_kernel_initrd()
{
    TMPDIR=$(mktemp -d)
    cd "$TMPDIR"
    mkdir lower upper workdir chroot initrd chroot/dev chroot/proc chroot/sys

    mount "$DISK"/casper/filesystem.squashfs lower
    mount -t overlay -o lowerdir="$TMPDIR"/lower,upperdir="$TMPDIR"/upper,workdir="$TMPDIR"/workdir none chroot
    mount -o bind /dev chroot/dev
    mount -o bind /proc chroot/proc
    mount -o bind /sys chroot/sys
    CLEANUP=2

    cd chroot
    apt-get download "$@"
    mkdir -p "$DISK"/debs/main
    cp -v linux*.deb "$DISK"/debs
    cd ..

    cat > chroot/generate_initrd.sh <<ENDLINE
#!/bin/bash

dpkg -i /*.deb
ENDLINE
    chmod +x chroot/generate_initrd.sh
    chroot chroot /generate_initrd.sh

    if [ -d upper/boot ]; then
        chmod 644 upper/boot/*
    fi
}

update_kernel_initrd ()
{
    release="$1"
    echo "INFO: using kernel release $release"
    ls -l upper/boot/
    [ -f "$DISK"/casper/vmlinuz.efi ] && cp -v "upper/boot/vmlinuz-$release" "$DISK"/casper/vmlinuz.efi
    [ -f "$DISK"/casper/vmlinuz ] && cp -v "upper/boot/vmlinuz-$release" "$DISK"/casper/vmlinuz

    cp "upper/boot/initrd.img-$release" .

    while file "initrd.img-$release" | grep cpio >/dev/null 2>&1; do
        offset=$(cpio -t <"initrd.img-$release" 2>&1 >/dev/null | awk '{print $1}')
        dd if="initrd.img-$release" of="initrd.img-$release.microcode" bs=512 count="$offset" >/dev/null 2>&1
        dd if="initrd.img-$release" of="initrd.img-$release.tmp" bs=512 skip="$offset" >/dev/null 2>&1
        mv "initrd.img-$release.tmp" "initrd.img-$release"
        cat <"initrd.img-$release.microcode" >>microcode.cpio
        rm "initrd.img-$release.microcode"
    done

    cd initrd
    if file "../initrd.img-$release" | grep -i gzip >/dev/null 2>&1; then
        zcat "../initrd.img-$release" | cpio --quiet -id
    elif file "../initrd.img-$release" | grep -i lzma >/dev/null 2>&1; then
        lzcat "../initrd.img-$release" | cpio --quiet -id
    elif file "../initrd.img-$release" | grep -i lz4 >/dev/null 2>&1; then
        lz4cat "../initrd.img-$release" | cpio --quiet -id
    else
        file "../initrd.img-$release"
        echo "WARNING: It doesn't deal with this format yet."
    fi
    rm "../initrd.img-$release"
    UUID="$(uuidgen -r)"
    echo "$UUID" > conf/uuid.conf
    find . 2>/dev/null | cpio --quiet -o -H newc | xz --verbose -9 --format=lzma >"../initrd.lz"
    cd ..

    if [ -f "microcode.cpio" ]; then
        CONCAT=("microcode.cpio" "initrd.lz")
    else
        CONCAT=("initrd.lz")
    fi

    [ -e "$DISK"/casper/initrd ] && cat "${CONCAT[@]}" > "$DISK"/casper/initrd
    [ -e "$DISK"/casper/initrd.lz ] && cat "${CONCAT[@]}" > "$DISK"/casper/initrd.lz
    echo "$UUID" > "$DISK"/.disk/casper-uuid
    CLEANUP=3
}

update_recovery_partition()
{
    [ -z "$1" ] && exit
    trap clean_up EXIT
    DISK=$(mktemp -d)
    mount "$1" "$DISK"
    CLEANUP=1
    cd "$DISK" || exit
    VERSION=$(file "$DISK/casper/vmlinuz" | grep -oP "version [^ ]*" | awk '{print $2}')
    case "$VERSION" in
        (4.4.0-*-generic)
            if dpkg --compare-versions "$VERSION" lt '4.4.0-184-generic'; then
                generate_kernel_initrd linux-headers-4.4.0-187 linux-headers-4.4.0-187-generic linux-image-4.4.0-187-generic linux-modules-4.4.0-187-generic linux-modules-extra-4.4.0-187-generic linux-base
                update_kernel_initrd 4.4.0-187-generic
                if [ ! -e "$DISK"/scripts/chroot-scripts/fish/99-install-kernel-4.4.0-187-generic.sh ]; then
                    cat > "$DISK"/scripts/chroot-scripts/fish/99-install-kernel-4.4.0-187-generic.sh <<ENDLINE
#!/bin/sh

apt-get install --yes --allow-unauthenticated linux-headers-4.4.0-187 linux-headers-4.4.0-187-generic linux-image-4.4.0-187-generic linux-modules-4.4.0-187-generic linux-modules-extra-4.4.0-187-generic
ENDLINE
                fi
            else
                echo "Linux kernel in recovery partition is new enough. No need to update it."
            fi
            ;;
        (4.8.0-*-generic|4.10.0-*-generic|4.11.0-*-generic|4.13.0-*-generic|4.13.0-*-oem)
            generate_kernel_initrd linux-headers-4.15.0-112 linux-headers-4.15.0-112-generic linux-image-4.15.0-112-generic linux-modules-4.15.0-112-generic linux-modules-extra-4.15.0-112-generic linux-base
            update_kernel_initrd 4.15.0-112-generic
            if [ ! -e "$DISK"/scripts/chroot-scripts/fish/99-install-kernel-4.15.0-112-generic.sh ]; then
                cat > "$DISK"/scripts/chroot-scripts/fish/99-install-kernel-4.15.0-112-generic.sh <<ENDLINE
#!/bin/sh

apt-get install --yes --allow-unauthenticated linux-headers-4.15.0-112 linux-headers-4.15.0-112-generic linux-image-4.15.0-112-generic linux-modules-4.15.0-112-generic linux-modules-extra-4.15.0-112-generic
ENDLINE
            fi
            ;;
        (4.15.0-*-generic)
            if dpkg --compare-versions "$VERSION" lt '4.15.0-101-generic'; then
                generate_kernel_initrd linux-headers-4.15.0-112 linux-headers-4.15.0-112-generic linux-image-4.15.0-112-generic linux-modules-4.15.0-112-generic linux-modules-extra-4.15.0-112-generic linux-base
                update_kernel_initrd 4.15.0-112-generic
                if [ ! -e "$DISK"/scripts/chroot-scripts/fish/99-install-kernel-4.15.0-112-generic.sh ]; then
                    cat > "$DISK"/scripts/chroot-scripts/fish/99-install-kernel-4.15.0-112-generic.sh <<ENDLINE
#!/bin/sh

apt-get install --yes --allow-unauthenticated linux-headers-4.15.0-112 linux-headers-4.15.0-112-generic linux-image-4.15.0-112-generic linux-modules-4.15.0-112-generic linux-modules-extra-4.15.0-112-generic
ENDLINE
                fi
            else
                echo "Linux kernel in recovery partition is new enough. No need to update it."
            fi
            ;;
    esac
}

case "$1" in
    (configure)
        if dpkg-query -W -f='${Status}\n' dell-recovery 2>&1 | grep "install ok installed" >/dev/null 2>&1; then
            TARGET=$(python3 -c "import Dell.recovery_common as magic; target=magic.find_partition(); print(target.decode('utf8')) if type(target) is bytes else print(target)")
        elif dpkg-query -W -f='${Status}\n' ubuntu-recovery 2>&1 | grep "install ok installed" >/dev/null 2>&1; then
            TARGET=$(python3 -c "import ubunturecovery.recovery_common as magic; target=magic.find_partition('PQSERVICE'); print(target.decode('utf8')) if type(target) is bytes else print(target)")
        else
            TARGET=""
        fi
        update_recovery_partition "$TARGET"
    ;;
esac

#DEBHELPER#
