#!/bin/bash
set -e

if [ -n "$D" ]; then
    set -x
fi

show_help() {
    echo "usage: repack-kernel <command> <opts>"
    echo ""
    echo "handle extraction of the kernel snap to a workspace directory, and later"
    echo "repacking it back to a snap."
    echo ""
    echo "   repack-kernel setup                        - setup system dependencies"
    echo "   repack-kernel extract <snap-file> <target> - extract under <target> workspace tree"
    echo "   repack-kernel prepare <target>             - prepare initramfs & kernel for repacking"
    echo "   repack-kernel pack <target>                - pack the kernel"
    echo "   repack-kernel cull-firmware <target>       - remove unnecessary firmware"
    echo "   repack-kernel cull-modules <target>        - remove unnecessary modules"
    echo ""
}

setup() {
    if [ "$UID" != "0" ]; then
        echo "run as root (only this command)"
        exit 1
    fi

    # carries ubuntu-core-initframfs
    add-apt-repository ppa:snappy-dev/image -y
    apt update
    
    if ! dpkg -s ubuntu-core-initramfs >/dev/null 2>&1; then
        if ! apt install -y ubuntu-core-initramfs; then
            echo "ubuntu-core-initramfs was not available in deb repository!"
            exit 1
        fi
    fi
}

get_kver() {
    local kerneldir="$1"
    find "$kerneldir" -maxdepth 1 -name "config-*" | grep -Po 'config-\K.*'
}

extract() {
    local snap_file="$1"
    local target="$2"

    if [ -z "$target" ] || [ -z "$snap_file" ]; then
        echo "repack-kernel: invalid arguments for extract"
        exit 1
    fi

    target=$(realpath "$target")

    mkdir -p "$target" "$target/work" "$target/backup"

    # kernel snap is huge, unpacking to current dir
    unsquashfs -d "$target/kernel" "$snap_file"

    kver="$(get_kver "$target/kernel")"

    # repack initrd magic, beware
    # assumptions: initrd is compressed with LZ4, cpio block size 512, microcode
    # at the beginning of initrd image

    # XXX: ideally we should unpack the initrd, replace snap-boostrap and
    # repack it using ubuntu-core-initramfs --skeleton=<unpacked> this does not
    # work and the rebuilt kernel.efi panics unable to start init, but we
    # still need the unpacked initrd to get the right kernel modules
    objcopy -j .initrd -O binary "$target"/kernel/kernel.efi "$target/work/initrd"

    # copy out the kernel image for create-efi command
    objcopy -j .linux -O binary "$target"/kernel/kernel.efi "$target/work/vmlinuz-$kver"

    cp -a "$target"/kernel/kernel.efi "$target/backup/"

    # this works on 20.04 but not on 18.04
    unmkinitramfs "$target"/work/initrd "$target"/work/unpacked-initrd
    
    # copy the unpacked initrd to use as the target skeleton
    cp -ar "$target/work/unpacked-initrd" "$target/skeleton"

    echo "prepared workspace at $target"
    echo "  kernel:              $target/kernel ($kver)"
    echo "  kernel.efi backup:   $target/backup/kernel.efi"
    echo "  temporary artifacts: $target/work"
    echo "  initramfs skeleton:  $target/skeleton"

}

prepare() {
    local target="$1"

    if [ -z "$target" ]; then
        echo "repack-kernel: missing target for prepare"
        exit 1
    fi

    target=$(realpath "$target")
    
    local kver
    kver="$(get_kver "$target/kernel")"

    (
        # all the skeleton edits go to a local copy of distro directory
        local skeletondir="$target/skeleton"

        cd "$target/work"
        # XXX: need to be careful to build an initrd using the right kernel
        # modules from the unpacked initrd, rather than the host which may be
        # running a different kernel
        (
            # accommodate assumptions about tree layout, use the unpacked initrd
            # to pick up the right modules
            cd unpacked-initrd/main
            ubuntu-core-initramfs create-initrd \
                                  --kernelver "$kver" \
                                  --skeleton "$skeletondir" \
                                  --feature main \
                                  --kerneldir "lib/modules/$kver" \
                                  --output "$target/work/repacked-initrd"
        )

        # assumes all files are named <name>-$kver
        ubuntu-core-initramfs create-efi \
                              --kernelver "$kver" \
                              --initrd repacked-initrd \
                              --kernel vmlinuz \
                              --output repacked-kernel.efi

        cp "repacked-kernel.efi-$kver" "$target/kernel/kernel.efi"

        # XXX: needed?
        chmod +x "$target/kernel/kernel.efi"
    )
}

cull_firmware() {
    local target="$1"

    if [ -z "$target" ]; then
        echo "repack-kernel: missing target for cull-firmware"
        exit 1
    fi
    
    # XXX: drop ~450MB+ of firmware which should not be needed in under qemu
    # or the cloud system
    rm -rf "$target"/kernel/firmware/*
}

cull_modules() {
    local target="$1"

    if [ -z "$target" ]; then
        echo "repack-kernel: missing target for cull-modules"
        exit 1
    fi

    target=$(realpath "$target")
    
    local kver
    kver="$(get_kver "$target/kernel")"

    (
        cd "$target/kernel"
        # drop unnecessary modules
        awk '{print $1}' <  /proc/modules  | sort > "$target/work/current-modules"
        #shellcheck disable=SC2044
        for m in $(find modules/ -name '*.ko'); do
            noko=$(basename "$m"); noko="${noko%.ko}"
            if echo "$noko" | grep -f "$target/work/current-modules" -q ; then
                echo "keeping $m - $noko"
            else
                rm -f "$m"
            fi
        done

        # depmod assumes that /lib/modules/$kver is under basepath
        mkdir -p fake/lib
        ln -s "$PWD/modules" fake/lib/modules
        depmod -b "$PWD/fake" -A -v "$kver"
        rm -rf fake
    )
}

pack() {
    local target="$1"

    if [ -z "$target" ]; then
        echo "repack-kernel: missing target for pack"
        exit 1
    fi
    snap pack "${@:2}" "$target/kernel" 
}

main() {
    if [ $# -eq 0 ]; then
        show_help
        exit 0
    fi

    local subcommand="$1"
    local action=
    while [ $# -gt 0 ]; do
        case "$1" in
            -h|--help)
                show_help
                exit 0
                ;;
            *)
                action=$(echo "$subcommand" | tr '-' '_')
                shift
                break
                ;;
        esac
    done

    if [ -z "$(declare -f "$action")" ]; then
        echo "repack-kernel: no such command: $subcommand" >&2
        show_help
        exit 1
    fi

    "$action" "$@"
}

main "$@"