ubuntu-kernel-cleanup.py

#!/usr/bin/python3
#
"""
  Remove intermediates and unused kernels from Ubuntu
  Please check /etc/ubuntu-kernel-cleanup.ini for tweaking & information

Usage:
  remove-ubuntu-kernels.py (--real-run | --dry-run)
  remove-ubuntu-kernels.py -h | --help

Options:
  -h --help       Show this screen
  --real-run      Real execution
  --dry-run       Dry execution
"""
import re
import os
import subprocess as sp
import configparser
from distutils.version import LooseVersion
import platform
import apt
try:
    from docopt import docopt
except ModuleNotFoundError as err:
    print("please install \'docopt\' module for python3")

try:
    from packaging.version import parse
except ModuleNotFoundError as err:
    print("please install \'packaging\' module for python3")


def get_packages_list(pkg_match):
    """
      run a shell command like this:
      dpkg --list linux-image-[1-9]* | awk '/^ii /&&/linux-image-/{print $2}'
      and create a sorted list of installed kernels.
      Or if we need the version only:
      dpkg-query -f '${Status} ${Version}\n' -W linux-image-[1-9]* | awk '/ok installed/{print $NF}'
    """
    installed = sp.Popen(
        "dpkg --list %s[1-9]* | awk \'/^ii /&&/%s/{print $2}\'" % (
            pkg_match, pkg_match),
        stdout=sp.PIPE, stderr=sp.PIPE, shell=True
    )
    installed_out, _ = installed.communicate()
    installed_list = installed_out.decode('utf-8').split('\n')
    # delete empty items
    kernels = [item for item in installed_list if item != '']
    kernel_versions = [item.replace(pkg_match, '') for item in kernels]
    # sort by version and unique
    return sorted(list(set(kernel_versions)), key=LooseVersion)


def remove_pkg(version_number, prefix, execution_type):
    """ remove given package """
    for suffix in KERNEL_SUFFIXES:
        pkg = "{}-{}{}".format(prefix, version_number, suffix)
        if execution_type == 'real':
            CACHE.update()
            CACHE.open()
        try:
            CACHE[pkg].is_installed
        except KeyError:
            pass
        else:
            if execution_type == 'real':
                CACHE[pkg].mark_delete(True, True)
                try:
                    CACHE.commit()
                except Exception as err:  # pylint: disable=W0703
                    print("Package removal failed [{err}]".format(
                        err=str(err)))
                finally:
                    CACHE.close()
            else:
                print('Would have removed {} (noop)'.format(pkg))


def get_version(full_pkg_name):
    """
      get version number from the full package name
      for instance linux-generic-5.8.0-34-generic returns 5.8.0-34
    """
    return re.sub(r"(^\D+-|-\D+)", r"", full_pkg_name)


if __name__ == '__main__':

    if os.geteuid() != 0:
        print("You need to have root privileges to run this script")
        print("Please try again, this time using 'sudo'. Exiting.")
        os.sys.exit()

    CFG_ONE = os.path.join(
        os.path.expanduser('~'), ".ubuntu-kernel-cleanup.ini")
    CFG_TWO = '/etc/ubuntu-kernel-cleanup.ini'
    if not os.path.isfile(CFG_ONE) and not os.path.isfile(CFG_TWO):
        print("you need to create either {} or {}.\nGiving up & waiting for better days to come")
        os.sys.exit()

    CONFIG = configparser.RawConfigParser()
    CONFIG.read([CFG_ONE, CFG_TWO])
    try:
        KERNEL_PREFIXES = CONFIG.get(
            'ubuntu-kernel-cleanup', 'kernel_prefixes').replace(' ', '').split(',')
    except configparser.NoOptionError:
        print('could not find an array option for "kernel_prefixes". Exiting')
        os.sys.exit()

    try:
        _KERNEL_SUFFIXES = CONFIG.get(
            'ubuntu-kernel-cleanup', 'kernel_suffixes').replace(' ', '').split(',')
    except configparser.NoOptionError:
        print('could not find an array option for "kernel_suffixes". Exiting')
        os.sys.exit()

    KERNEL_SUFFIXES = [''] + ['-{}'.format(x) for x in _KERNEL_SUFFIXES]

    try:
        COUNT = CONFIG.getint('ubuntu-kernel-cleanup', 'count')
    except configparser.NoOptionError:
        print('could not find an option for "count". Defaulting to 2')
        COUNT = 2

    if COUNT < 1:
        print("--count must greater or equal to 1")
        os.sys.exit()

    EXECUTION = 'dry'
    ARGS = docopt(__doc__)
    if ARGS.get('--real-run'):
        EXECUTION = 'real'

    running_kernel = get_version(platform.uname().release)

    # load packages cache
    CACHE = apt.cache.Cache()

    kern_list = get_packages_list('linux-image-')

    for kern_count in list(range(COUNT)):
        # do we still have kernels in the list?
        if len(kern_list) > 1:
            latest_kernel = kern_list[-1]
            if parse(running_kernel) < parse(latest_kernel):
                # keep only items not containing the latest version
                sanitized_kernel_list = [
                    item for item in kern_list if latest_kernel != item]
            else:
                # we are running the latest kernel
                sanitized_kernel_list = kern_list
                del kern_list[-1]

    try:
        sanitized_kernel_list
    except NameError:
        pass
    else:
        # delete running kernel package from list
        purged_list = [
            item for item in sanitized_kernel_list if running_kernel not in item]

        # uninstall packages from list
        for kernel_pkg in KERNEL_PREFIXES:
            for kernel_version in purged_list:
                remove_pkg(kernel_version, kernel_pkg, EXECUTION)

    CACHE.close()  # it returns always true... keeps closing :)

    # purge empty packages
    if EXECUTION == 'real':
        CMD = "dpkg-query -f \'${Package} ${Status}\\n\' -W linux-* | awk  \'/deinstall ok/{print $1}\' | xargs apt-get purge -y"  # pylint: disable=C0301
    else:
        CMD = "dpkg-query -f \'${Package} ${Status}\\n\' -W linux-* | awk  \'/deinstall ok/{print $1}\' | xargs apt-get purge --assume-no"  # pylint: disable=C0301

    purge = sp.Popen(CMD, stdout=sp.PIPE, stderr=sp.PIPE, shell=True)
    purge_out, _ = purge.communicate()
    print("\ntrying to purge empty kernel packages:\n{}".format(
        purge_out.decode('utf-8')))