#!/usr/bin/python3
from __future__ import division
import posixpath
import re
import subprocess
import sys


class Version(object):
    """Crude version abstraction for systems without distutils.version"""

    def __init__(self, version_str):
        self.version_str = version_str

    def __str__(self):
        return self.version_str

    def __repr__(self):
        return 'Version(\'%s\')' % self.version_str

    def components(self):
        # Split the version into components, by word boundaries, or a
        # change between numbers and non-numbers and vice-versa. re.split()
        # does not work on zero-length delimiters, so we have to use a
        # sub+split.
        return re.sub(r'\b|(?<=\d)(?=\D)|(?<=\D)(?=\d)', '\n',
                      self.version_str).split('\n')

    def __eq__(self, other):
        return self.version_str == other.version_str

    def __gt__(self, other):
        def num_gt(str_, str_other):
            m = re.search(r'^\d+', str_)
            m_other = re.search(r'^\d+', str_other)
            if m and m_other:
                return int(m.group(0)) > int(m_other.group(0))
            else:
                return str_ > str_other

        for self_c, other_c in zip(self.components(), other.components()):
            if self_c == other_c:
                continue
            else:
                return num_gt(self_c, other_c)
        return False

    # Note: not using functools.total_ordering to support Python 2.6

    def __lt__(self, other):
        return not (self == other) and not (self > other)


def proc_version():
    """Return the content of /proc/version"""

    proc_version = None

    # Not using a with statement here, to support Python 2.4
    v = open('/proc/version', 'r')
    try:
        proc_version = v.read()
    finally:
        v.close()
    return proc_version


def running_kernel_version():
    proc_version_ = proc_version()

    if re.search(r'Debian', proc_version_):
        # Remove gcc version first
        proc_version_ = re.sub(r'\(gcc.*?\)', '', proc_version_)

        # Then look for the Debian kernel version
        m = re.search(r'((?<=Debian )(\S+)|(?<=PVE )(\S+))', proc_version_)
        if m:
            version_str = m.group(1).strip('()')
            version = clean_kernel_version(version_str)
            return version

    else:
        m = re.search(r'(?<=Linux version )(\S+)', proc_version_)
        if m:
            version = clean_kernel_version(m.group(0))
            return version


def is_debian():
    return posixpath.exists('/etc/debian_version')


def is_fedora():
    return posixpath.exists('/etc/fedora-release')


def is_redhat():
    return posixpath.exists('/etc/redhat-release')


def installed_kernel_versions():
    """
    Debian-based kernels come in different flavors.

    On some of them, like Proxmox' PVE kernels, special care has to be taken
    when selecting the right packages to run `dpkg-query` on.
    """
    if is_debian():
        return \
          installed_kernel_versions_debian(pkgname='linux-image') \
          or installed_kernel_versions_debian(pkgname='pve-kernel',
                                              pkgquery='pve-kernel*-pve')
    if is_fedora() or is_redhat():
        return installed_kernel_versions_fedora()
    return [None]


def check_output(cmd):
    """Emulate subprocess.check_output for ancient Python versions"""
    return subprocess.Popen(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    ).communicate()[0]


def installed_kernel_versions_debian(pkgname='linux-image', pkgquery=None):
    if not pkgquery:
        pkgquery = '%s*' % pkgname
    dpkg_out = check_output(
        ['dpkg-query',
         '--show', '--showformat', '${Package} ${Version}\n',
         pkgquery])
    dpkg_out = dpkg_out.decode('ascii', 'ignore').strip()

    versions = dpkg_out.split('\n')
    versions = [v for v in versions if re.search(r'^%s-\d.* \S+$' % pkgname, v)]
    versions = [clean_kernel_version(v.split(' ')[1]) for v in versions]

    return versions


def installed_kernel_versions_fedora():
    rpm_out = check_output(
        ['rpm', '--queryformat=%{VERSION}-%{RELEASE}\n', '-q', 'kernel'])
    rpm_out = rpm_out.decode('ascii', 'ignore').strip()

    versions = rpm_out.split('\n')
    versions = [clean_kernel_version(v) for v in versions]

    return versions


def installed_kernel_version():
    installed = sorted(installed_kernel_versions())
    if not installed:
        print('KERNEL UNKNOWN - unable to determine installed kernel version')
        sys.exit(3)
    return installed[-1]


def clean_kernel_version(version):
    # arch
    version = re.sub(r'\.(x86_64|i[3-6]86)', '', version)
    # Fedora release
    version = re.sub(r'\.fc\d+', '', version)
    # RHEL release
    version = re.sub(r'\.el[\d\._]+$', '', version)

    return Version(version)


def main():
    if len(sys.argv) > 1:
        print('This plugin no longer takes the expected kernel version as' +
              'an argument')
        sys.exit(3)

    running = running_kernel_version()
    installed = installed_kernel_version()

    if running == installed:
        print('KERNEL OK - running version %s' % running)
        sys.exit(0)
    else:
        print('KERNEL WARNING - running version %s, installed: %s' %
              (running, installed))
        sys.exit(1)


if __name__ == '__main__':
    main()