#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
This script will delete all old non used kernels.
Tested on Ubuntu 12.04 and 14.04.
"""
import sys
import subprocess
import re


RED = '\033[91m'
GREEN = '\033[92m'
YELLOW = '\033[93m'
BLUE = '\033[94m'
PURPLE = '\033[95m'
TEAL = '\033[96m'
DEFAULT = '\033[0m'


def log(text, error=False):
    fo = sys.stderr if error else sys.stdout
    print(text, file=fo)
    fo.flush()


def clean_kernels():
    # check that the user can use sudo
    p = subprocess.Popen('sudo echo ok', stdin=sys.stdin, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    p.communicate()
    if p.returncode != 0:
        log(RED + 'Failed to get sudo right.' + DEFAULT, error=True)
        return 1
    # list installed kernels
    p = subprocess.Popen('sudo dpkg --get-selections | grep linux', stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    out, err = p.communicate()
    out = str(out, 'utf-8') if out else ''
    err = str(err, 'utf-8') if err else ''
    if p.returncode != 0:
        log(RED + 'Failed to list kernels.' + DEFAULT + '\nOut: ' + out + '\nErr: ' + err, error=True)
        return 1
    kernels = list()
    for line in out.split('\n'):
        m = re.match(r'linux-(headers|image)([\w-]*)-([\d\.-]+)', line)
        if m:
            version = m.groups()[2]
            if version[-1] == '-':
                version = version[:-1]
            kernels.append((line.split('\t')[0], version))
    log('Installed kernels:\n\t' + '\n\t'.join([n for n, v in kernels]))
    # get current kernel
    p = subprocess.Popen('uname -a', stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    out, err = p.communicate()
    out = str(out, 'utf-8') if out else ''
    err = str(err, 'utf-8') if err else ''
    if p.returncode != 0:
        log(RED + 'Failed to get current kernel.' + DEFAULT + '\nOut: ' + out + '\nErr: ' + err, error=True)
        return 1
    try:
        current = out.split(' ')[2]
        if current.endswith('-generic'):
            current = current[:-8]
    except Exception as e:
        log(RED + 'Failed to get current kernel.' + DEFAULT + '\nError: ' + e, error=True)
        return 1
    log('Current kernel is: ' + PURPLE + current + DEFAULT)
    # get kernel packages to purge
    to_purge = list()
    current_t = current.replace('-', '.').split('.')
    latest = current_t
    for name, version in kernels:
        version_t = version.replace('-', '.').split('.')
        if version_t > latest:
            latest = version_t
        elif version_t < latest and version != current:
            to_purge.append(name)
    if not to_purge:
        log('No kernel package to purge.')
        return 0
    log('Kernel packages to purge:\n\t' + '\n\t'.join(to_purge))
    try:
        r = input('Do you confirm the deletion of these packages? [y]/n ')
    except (KeyboardInterrupt, EOFError):
        log('')
    else:
        if not r or r in ('y', 'yes'):
            log(PURPLE + 'Starting purge...' + DEFAULT)
            cmd = ['sudo', 'apt-get', '--yes', 'purge']
            cmd.extend(to_purge)
            p = subprocess.Popen(cmd, stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stderr, shell=False)
            p.communicate()
            sys.stdout.flush()
            sys.stderr.flush()
            if p.returncode != 0:
                log(RED + 'Failed to purge kernel packages.' + DEFAULT, error=True)
                return 1
            log(GREEN + 'Packages purged.' + DEFAULT)
    return 0


if __name__ == '__main__':
    sys.exit(clean_kernels())