#!/bin/bash
# SPDX-License-Identifier: GPL-3.0-or-later
# PLONK - PLONK Leaves Only Needed Kernels
#
# This script removes old Debian kernels
# It also purges old kernel packages left in "rc" state
#
# In short: old kernels go plonk.
#
# Author: bisco <bisco__AT__younerd__DOT_org>
# Script inspired by the following cyberciti post:
# https://www.cyberciti.biz/faq/debian-ubuntu-linux-delete-old-kernel-images-command/

set -euo pipefail

### Variables

APT_GET="$(command -v apt-get || true)"
DPKG_QUERY="$(command -v dpkg-query || true)"

KERNELVER="$(uname -r)"
KEEP_KERNELS="2"

ACTION=""
ASSUME_YES="no"

KERNELS_TO_KEEP=()
ABIS_TO_KEEP=()

OLD_KERNELS=()
RC_KERNELS=()
PKGS_TO_PURGE=()

### End Variables


### Code.
### Please don't edit below unless if you know what you're doing.


### Check if something went wrong
check_error(){
    echo "====> ERROR!"
    echo "====> Script failed on line ${BASH_LINENO[0]}"
    exit 1
}


### Run check_error function on errors
trap check_error ERR


### Show help and exits
usage(){
    echo "PLONK - PLONK Leaves Only Needed Kernels"
    echo ""
    echo "Old kernels go plonk."
    echo ""
    echo "Usage: $0 [option]"
    echo ""
    echo "-l | --list      : list old kernel packages"
    echo "-n | --dry-run   : simulate old kernel packages removal"
    echo "-r | --remove    : remove old kernel packages"
    echo "-k | --keep NUM  : number of kernels to keep. Default: ${KEEP_KERNELS}"
    echo "-y | --yes       : assume yes when removing packages"
    echo "-h | --help      : show this help and exits"
    echo ""
    echo "Examples:"
    echo "$0 --list"
    echo "$0 --dry-run"
    echo "sudo $0 --remove"
    echo "sudo $0 --remove --keep 3"
    echo ""
}


### Check if all binaries are properly installed on target system
check_binary(){
    if [ "${APT_GET}" == "" ] || [ "${DPKG_QUERY}" == "" ]; then
        echo "[FAIL] Some binaries not found."
        echo "[INFO] Please make sure you have: apt-get, dpkg-query"
        exit 1
    else
        echo "[OK] All binaries are properly installed. Let's go on!"
    fi
}


### Check if the script has been run by root when removing packages
check_root(){
    if [ "${ACTION}" != "remove" ]; then
        return
    fi

    if [ "$(id -u)" == "0" ]; then
        echo "[OK] Check root permissions"
    else
        echo "[FAIL] Check root permissions"
        echo "[INFO] Run this script as root or sudo"
        exit 1
    fi
}


### Extract kernel version from package name
pkg_to_kernel_version(){
    local PKG="${1%%:*}"

    case "${PKG}" in
        linux-image-unsigned-*)
            echo "${PKG#linux-image-unsigned-}"
            ;;
        linux-image-*)
            echo "${PKG#linux-image-}"
            ;;
        linux-headers-*)
            echo "${PKG#linux-headers-}"
            ;;
        *)
            return 1
            ;;
    esac
}


### Extract Debian kernel ABI from kernel package version
kernel_abi_from_version(){
    local VERSION="$1"

    case "${VERSION}" in
        *-cloud-amd64)
            echo "${VERSION%-cloud-amd64}"
            ;;
        *-rt-amd64)
            echo "${VERSION%-rt-amd64}"
            ;;
        *-amd64)
            echo "${VERSION%-amd64}"
            ;;
        *-686-pae)
            echo "${VERSION%-686-pae}"
            ;;
        *-686)
            echo "${VERSION%-686}"
            ;;
        *-cloud-arm64)
            echo "${VERSION%-cloud-arm64}"
            ;;
        *-rt-arm64)
            echo "${VERSION%-rt-arm64}"
            ;;
        *-arm64)
            echo "${VERSION%-arm64}"
            ;;
        *-common)
            echo "${VERSION%-common}"
            ;;
        *)
            echo "${VERSION%-*}"
            ;;
    esac
}


### List installed Debian kernel packages
list_installed_kernel_packages(){
    {
        "${DPKG_QUERY}" -W -f='${db:Status-Abbrev}\t${binary:Package}\n' \
            'linux-image-[0-9]*' \
            'linux-image-unsigned-[0-9]*' \
            'linux-headers-[0-9]*' \
            2>/dev/null || true
    } | awk '$1 == "ii" { print $2 }' | sort -u
}


### List residual Debian kernel packages in rc state
list_rc_kernel_packages(){
    {
        "${DPKG_QUERY}" -W -f='${db:Status-Abbrev}\t${binary:Package}\n' \
            'linux-image-[0-9]*' \
            'linux-image-unsigned-[0-9]*' \
            'linux-headers-[0-9]*' \
            2>/dev/null || true
    } | awk '$1 == "rc" { print $2 }' | sort -u
}


### List installed Debian kernel image versions
list_installed_kernel_versions(){
    {
        "${DPKG_QUERY}" -W -f='${db:Status-Abbrev}\t${binary:Package}\n' \
            'linux-image-[0-9]*' \
            'linux-image-unsigned-[0-9]*' \
            2>/dev/null || true
    } | awk '$1 == "ii" { print $2 }' \
      | while read -r PKG; do
            pkg_to_kernel_version "${PKG}"
        done \
      | sort -V -u
}


### Build the list of kernels that should be kept
build_keep_list(){
    local INSTALLED_KERNELS=()
    local VERSION=""
    local ABI=""
    local I=""

    mapfile -t INSTALLED_KERNELS < <(list_installed_kernel_versions)

    KERNELS_TO_KEEP=("${KERNELVER}")

    for (( I=${#INSTALLED_KERNELS[@]}-1; I>=0 && ${#KERNELS_TO_KEEP[@]}<KEEP_KERNELS; I-- )); do
        VERSION="${INSTALLED_KERNELS[$I]}"

        if [ "${VERSION}" == "${KERNELVER}" ]; then
            continue
        fi

        KERNELS_TO_KEEP+=("${VERSION}")
    done

    ABIS_TO_KEEP=()

    for VERSION in "${KERNELS_TO_KEEP[@]}"; do
        ABI="$(kernel_abi_from_version "${VERSION}")"
        ABIS_TO_KEEP+=("${ABI}")
    done
}


### Check if a package belongs to a kernel that should be kept
is_kernel_to_keep(){
    local PKG="$1"
    local VERSION=""
    local ABI=""
    local KEEP=""

    VERSION="$(pkg_to_kernel_version "${PKG}")"
    ABI="$(kernel_abi_from_version "${VERSION}")"

    for KEEP in "${KERNELS_TO_KEEP[@]}"; do
        if [ "${VERSION}" == "${KEEP}" ]; then
            return 0
        fi
    done

    for KEEP in "${ABIS_TO_KEEP[@]}"; do
        if [ "${ABI}" == "${KEEP}" ]; then
            return 0
        fi
    done

    return 1
}


### List old installed kernel packages
list_old_installed_kernels(){
    local PKG=""

    while read -r PKG; do
        if [ "${PKG}" == "" ]; then
            continue
        fi

        if ! is_kernel_to_keep "${PKG}"; then
            echo "${PKG}"
        fi
    done < <(list_installed_kernel_packages) | sort -u
}


### Build final package list
build_purge_list(){
    mapfile -t OLD_KERNELS < <(list_old_installed_kernels)
    mapfile -t RC_KERNELS < <(list_rc_kernel_packages)

    PKGS_TO_PURGE=("${OLD_KERNELS[@]}" "${RC_KERNELS[@]}")
}


### Print script summary
print_summary(){
    echo ""
    echo "Running kernel : ${KERNELVER}"
    echo "Keeping kernels: ${KERNELS_TO_KEEP[*]}"
    echo "Keeping ABI    : ${ABIS_TO_KEEP[*]}"
    echo ""

    if [ "${#OLD_KERNELS[@]}" -gt 0 ]; then
        echo "Old installed kernel packages:"
        printf '  %s\n' "${OLD_KERNELS[@]}"
    else
        echo "[OK] No old installed kernel packages found."
    fi

    echo ""

    if [ "${#RC_KERNELS[@]}" -gt 0 ]; then
        echo "Residual rc kernel packages:"
        printf '  %s\n' "${RC_KERNELS[@]}"
    else
        echo "[OK] No residual rc kernel packages found."
    fi

    echo ""
}


### Simulate package purge
dry_run(){
    if [ "${#PKGS_TO_PURGE[@]}" -eq 0 ]; then
        echo "[OK] Nothing to purge."
        exit 0
    fi

    "${APT_GET}" -s purge "${PKGS_TO_PURGE[@]}"
}


### Remove old kernel packages
remove_kernels(){
    local ANSWER=""

    if [ "${#PKGS_TO_PURGE[@]}" -eq 0 ]; then
        echo "[OK] Nothing to purge."
        exit 0
    fi

    if [ "${ASSUME_YES}" != "yes" ]; then
        read -r -p "Proceed with apt-get purge? [y/N] " ANSWER

        case "${ANSWER}" in
            y|Y|yes|YES)
                ;;
            *)
                echo "Aborted."
                exit 0
                ;;
        esac
    fi

    if [ "${ASSUME_YES}" == "yes" ]; then
        "${APT_GET}" -y purge "${PKGS_TO_PURGE[@]}"
    else
        "${APT_GET}" purge "${PKGS_TO_PURGE[@]}"
    fi
    
    echo ""
    echo "Old kernels have been successfully plonked."
}


### Parse options
while [ "$#" -gt 0 ]; do
    case "$1" in
        -l | --list)
            ACTION="list"
            ;;
        -n | --dry-run)
            ACTION="dry-run"
            ;;
        -r | --remove)
            ACTION="remove"
            ;;
        -k | --keep)
            shift

            if ! [[ "${1:-}" =~ ^[0-9]+$ ]]; then
                echo "[FAIL] --keep requires a number"
                exit 1
            fi

            if [ "$1" -lt 1 ]; then
                echo "[FAIL] --keep requires a number greater than zero"
                exit 1
            fi

            KEEP_KERNELS="$1"
            ;;
        -y | --yes)
            ASSUME_YES="yes"
            ;;
        -h | --help)
            usage
            exit 0
            ;;
        *)
            usage
            exit 1
            ;;
    esac

    shift
done


### Main
if [ "${ACTION}" == "" ]; then
    usage
    exit 1
fi

check_binary
check_root
build_keep_list
build_purge_list
print_summary

case "${ACTION}" in
    list)
        exit 0
        ;;
    dry-run)
        dry_run
        ;;
    remove)
        remove_kernels
        ;;
esac