#!/usr/bin/python
# -*- coding: utf-8 -*-
#
#    Logaricheck - a nagios plugin to check disk, memory and swap usage with
#                  intelligent thresolds.
#
#    Copyright (C) 2006-2010  Révolution Linux (http://www.revolutionlinux.com)
#              (C) 2006-2010  Guillaume Pratte <guillaume.pratte@revolutionlinux.com>
#              (C) 2010       Michael Jeanson <mjeanson@rlnx.com>
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License along
#    with this program; if not, write to the Free Software Foundation, Inc.,
#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.


RET_OK = 0
RET_WARNING = 1
RET_CRITICAL = 2
RET_UNKNOWN = 3

from optparse import OptionParser
import traceback, sys, math, os, statvfs
import ConfigParser

#### Filesystem ####

class FileSystemInfo:
    def __init__(self, device, mount_path, fstype):
        self.device = device
        self.mount_path = mount_path
        self.type = fstype
        self.total_size = 0
        self.available_size = 0
        # some filesystems such as ext3 reserve blocks for root,
        # so this size might be superior than the number of blocks
        # avaible to non-root users
        self.available_size_for_root = 0
        self.total_inodes = 0
        self.available_inodes = 0
        self.available_inodes_for_root = 0

        # Get disk usage
        statvfs_result = os.statvfs(self.mount_path)

        ## Blocks
        # Bytes per block
        block_size = statvfs_result[statvfs.F_BSIZE]

        # Free space available to non-root users
        self.available_size = statvfs_result[statvfs.F_BAVAIL] * block_size

        # Free space available to root
        self.available_size_for_root = statvfs_result[statvfs.F_BFREE] * block_size

        # Total space
        self.total_size = statvfs_result[statvfs.F_BLOCKS] * block_size

        ## Inodes
        # Free inodes available to non-root user.
        self.available_inodes = statvfs_result[statvfs.F_FAVAIL]

        # Free inodes available to root.
        self.available_inodes_for_root = statvfs_result[statvfs.F_FFREE]

        # Total number of file nodes.
        self.total_inodes = statvfs_result[statvfs.F_FILES]

    def __str__(self):
        return "%s %s %s total=%s avail=%s for_root=%s totali=%s availi=%s for_rooti=%s" % \
            (self.device, self.mount_path, self.type, self.total_size, self.available_size, \
             self.available_size_for_root, self.total_inodes, self.available_inodes, self.available_inodes_for_root)

def get_all_filesystem_info():
    # Open /etc/mtab and extract the filesystem list
    f = open("/etc/mtab")
    lines = f.readlines()
    f.close()

    fs_list = []

    # For each filesystem
    for line in lines:
        splitted_line = line.split()
        # (device, mount_point, fstype)
        try:
            fs = FileSystemInfo(splitted_line[0], splitted_line[1], splitted_line[2])
            # don't report special filesystems (such as /proc, ...), 
            # which have a total size of 0
            if fs.total_size != 0:
                fs_list.append(fs)
        except OSError:
            # Ignore cases where the user can't access the file system (fuse)
            pass
    return fs_list

class AbstractFileSystemUsage:
    def __init__(self, filesysteminfo):
        self.fsinfo = filesysteminfo
        self.type = None

    def get_total(self):
        raise NotImplementedError
        
    def get_available(self):
        raise NotImplementedError
    
    def get_available_for_root(self):
        raise NotImplementedError

    def get_real_usage(self):
        """ Total minus available for root """
        return self.get_total() - self.get_available_for_root()

    def get_non_root_usage(self):
        """ Total minus available for non-root users """
        return self.get_total() - self.get_available()

    def get_real_usage_percent(self):
        return (100.0 * self.get_real_usage()) / self.get_total()

    def get_non_root_usage_percent(self):
        return (100.0 * self.get_non_root_usage()) / self.get_total()

    def get_available_reserved_for_root(self):
        return self.get_available_for_root() - self.get_available()

class FileSystemSpaceUsage(AbstractFileSystemUsage):
    def get_total(self):
        return self.fsinfo.total_size
        
    def get_available(self):
        return self.fsinfo.available_size
    
    def get_available_for_root(self):
        return self.fsinfo.available_size_for_root
    
class FileSystemInodeUsage(AbstractFileSystemUsage):
    def get_total(self):
        return self.fsinfo.total_inodes
        
    def get_available(self):
        return self.fsinfo.available_inodes
    
    def get_available_for_root(self):
        return self.fsinfo.available_inodes_for_root

#### /Filesystem ####

#### meminfo ####

class RawMemoryInfo:
    def __init__(self):
        # Get memory usage information, according to /proc/meminfo
        f = open('/proc/meminfo', 'r')
        self.ram_buffers = 0
        self.ram_cached = 0
        self.ram_total = 0
        self.ram_free = 0
        self.swap_total = 0
        self.swap_free = 0

        for line in f.readlines():
            tokens = line.split()
            label = tokens[0]
            size = tokens[1]
            try:
                size = int(size)
            except ValueError:
                # The line is an header, skip it.
                continue

            # We have to approximate kb to bytes...
            size = size * 1024
            if label == 'MemTotal:':
                self.ram_total = size
            elif label == 'MemFree:':
                self.ram_free = size
            elif label == 'Buffers:':
                self.ram_buffers = size
            elif label == 'Cached:':
                self.ram_cached = size
            elif label == 'SwapTotal:':
                self.swap_total = size 
            elif label == 'SwapFree:':
                self.swap_free = size
        f.close()

class AbstractMemoryUsage:
    def __init__(self):
        self.total = 0
        self.used = 0

    def __str__(self):
        return '%s used, %s free' % (self.used, self.get_free())

    def get_free(self):
        return self.total - self.used

    def get_used_percent(self):
        return (float(self.used) / self.total) * 100
    
class RAMMemoryUsage(AbstractMemoryUsage):
    def __init__(self):
        AbstractMemoryUsage.__init__(self)

        meminfo = RawMemoryInfo()
        self.total = meminfo.ram_total
        self.used = meminfo.ram_total - meminfo.ram_free - meminfo.ram_buffers - meminfo.ram_cached

class SwapMemoryUsage(AbstractMemoryUsage):
    def __init__(self):
        AbstractMemoryUsage.__init__(self)

        meminfo = RawMemoryInfo()
        self.total = meminfo.swap_total
        self.used = self.total - meminfo.swap_free

#### /meminfo ####

def exit_ok(msg=''):
    print 'OK', msg
    sys.exit(RET_OK)

def exit_warning(msg=''):
    print 'WARNING', msg
    sys.exit(RET_WARNING)

def exit_critical(msg=''):
    print 'CRITICAL', msg
    sys.exit(RET_CRITICAL)

def round_float(val):
    return int(round(val))

def round_percent(val):
    return round_float(100 * val)

def to_mb(val):
    return round_float(val / (1024.0 * 1024.0))

def get_thresold(value, sensitivity, agressivity):
    # Invert logarithmic function to calculate appropriate thresold according to
    # total (disk|mem|swap|inode) space
    ln = math.log(value+10, 10)
    val = (sensitivity / 100.0) / (ln ** agressivity)
    if val>1:
        val = 0.99999
        
    #print "x=%s ; sens=%s ; agres=%s ; res=%s" % (value, sensitivity, agressivity, val)
    return val

def get_thresold_mb(value, sensitivity, agressivity):
    return get_thresold(to_mb(value), sensitivity, agressivity)

class ConfigItem:
    def __init__(self, sensitivity_warning, sensitivity_critical, agressivity, display_unit, display_divider):
        self.warning = sensitivity_warning
        self.critical = sensitivity_critical
        self.agressivity = agressivity
        self.display_unit = display_unit
        self.display_divider = display_divider
        
    def __str__(self):
        return "<ConfigItem : warn=%s; crit=%s; agres=%s; disp_unit=%s; disp_div=%s>" % \
            (self.warning, self.critical, self.agressivity, self.display_unit, self.display_divider)

class Configuration:
    def __init__(self):
        config_loaded = False

        confparser = ConfigParser.SafeConfigParser()
        if confparser.read("/etc/logaricheck/logaricheck.conf"):
            config_loaded = True

        # Default values, if the .conf file is not present.
        # sensitivity_warning, sensitivity_critical, agressivity, display_unit, display_divider
        self.memory = ConfigItem(250, 100, 2, "MB", 1024*1024)  
        self.swap = ConfigItem(300, 150, 2, "MB", 1024*1024)
        self.disk = ConfigItem(150, 50, 2, "GB", 1024*1024*1024)
        self.disk_inode = ConfigItem(40, 25, 1, "M", 1000*1000)
        self.fstype_ignore_list = ['iso9660','tmpfs'] # Default list

        if config_loaded:
            # FIXME : Refact needed, we are doing thrice the same thing
            for confname in ('memory', 'swap', 'disk', 'disk_inode'):
                confitem = getattr(self, confname)
                # Integer options
                for intopt in ("warning", "critical", "display_divider"):
                    try:
                        setattr(confitem, intopt, confparser.getint(confname, intopt))
                    except (ConfigParser.NoOptionError, ConfigParser.NoSectionError):
                        # If not set in config file, use default
                        pass
                # Float options
                for fltopt in ("agressivity",):
                    try:
                        setattr(confitem, fltopt, confparser.getfloat(confname, fltopt))
                    except (ConfigParser.NoOptionError, ConfigParser.NoSectionError):
                        # If not set in config file, use default
                        pass
                # String options
                for stropt in ("display_unit",):
                    try:
                        setattr(confitem, stropt, confparser.get(confname, stropt))
                    except (ConfigParser.NoOptionError, ConfigParser.NoSectionError):
                        # If not set in config file, use default
                        pass

            if confparser.has_section("filesystem") and confparser.has_option("filesystem", "ignore_list"):
                # Create a list separating the items with ',' and ignoring spaces
                self.fstype_ignore_list = confparser.get("filesystem", "ignore_list").replace(" ","").split(",")

class DiskCheck:
    def check_disk_usage(self, config):
        return self.check_usage(FileSystemSpaceUsage, config.disk, config.fstype_ignore_list)

    def check_disk_inode_usage(self, config):
        return self.check_usage(FileSystemInodeUsage, config.disk_inode, config.fstype_ignore_list)
        
    # Common code for disk and disk_inode validations
    def check_usage(self, FileSystemUsageClass, configitem, fstype_ignore_list):
        ok_msg = ""
        warning_list = []
        critical_list = []
        for fsinfo in get_all_filesystem_info():
            # Ignore certain filesystems
            if fsinfo.type in fstype_ignore_list:
                continue
            
            fsusage = FileSystemUsageClass(fsinfo)

            total = fsusage.get_total()
            used = fsusage.get_real_usage()
            free = fsusage.get_available()

            warning_free_thresold = get_thresold_mb(total, configitem.warning, configitem.agressivity)
            critical_free_thresold = get_thresold_mb(total, configitem.critical, configitem.agressivity)

            free_percent = free / float(total)

            message = "%s: %s%%  %.2f %s / %.2f %s used (warn=%s%% crit=%s%%)" % (fsinfo.mount_path,
                round_percent(1 - free_percent), 
                used / float(configitem.display_divider), configitem.display_unit,
                total/ float(configitem.display_divider), configitem.display_unit,
                round_percent(1 - warning_free_thresold), round_percent(1 - critical_free_thresold))

            if free_percent < critical_free_thresold:
                critical_list.append(message)
            elif free_percent < warning_free_thresold:
                warning_list.append(message)
            if fsinfo.mount_path == "/":
                ok_msg = message

        to_print = ''

        if critical_list:
            to_print += "CRITICAL<br>"
            for item in critical_list:
                to_print += item + '<br>'
        if warning_list:
            to_print += "WARNING<br>"
            for item in warning_list:
                to_print += item + '<br>'
        if to_print != '':
            print to_print
        if critical_list:
            sys.exit(RET_CRITICAL)
        if warning_list:
            sys.exit(RET_WARNING)
        exit_ok(ok_msg)
    
class MemoryCheck:
    def check_memory_usage(self, config):
        ramusage = RAMMemoryUsage()
        return self.check_usage(ramusage, config.memory)

    def check_swap_usage(self, config):
        swapusage = SwapMemoryUsage()
        if swapusage.total <= 0:
            #FIXME: Add configurable exit status when no swap is enabled.
            exit_ok("No swap enabled.")
        return self.check_usage(swapusage, config.swap)

    # Common code for mem and swap validation
    def check_usage(self, memusage, configitem):
        free = memusage.total - memusage.used
        free_percent = free / float(memusage.total)
        used_percent = 1 - free_percent

        warning_free_thresold = get_thresold_mb(memusage.total, configitem.warning, configitem.agressivity)
        critical_free_thresold = get_thresold_mb(memusage.total, configitem.critical, configitem.agressivity)

        message = "%s%%  %.1f %s / %.1f %s used (warn=%s%% crit=%s%%)" % \
            (round_percent(used_percent),
            memusage.used / float(configitem.display_divider), configitem.display_unit,
            memusage.total / float(configitem.display_divider), configitem.display_unit,
            round_percent(1 - warning_free_thresold),
            round_percent(1 - critical_free_thresold))

        if free_percent < critical_free_thresold:
            exit_critical(message)
        if free_percent < warning_free_thresold:
            exit_warning(message)
        exit_ok(message)

class PrintCheck:
    def print_thresolds(self, config):
        """ Print thresolds for debug purpose """
        print "The numbers represents the thresolds to raise an alert when the free disk/memory/swap\n\
is lower than the calculated percentage. This percentage varies according to\n\
the total capacity of the disk/memory/swap."
        
        values = (256, 512, 1024, 2048, 4096, 8192, 16384, 32768)
        self.print_thresold("Memory", values, "MB", 1024 * 1024, config.memory)
        self.print_thresold("Swap", values, "MB", 1024 * 1024, config.swap)

        values = (0.001, 0.01, 0.1, 0.5, 1, 5, 10, 50, 100, 500, 1000, 5000, 10000, 50000, 100000)
        self.print_thresold("Disk space", values, "GB", 1024 * 1024 * 1024, config.disk)

        values = (0.5, 1, 5, 10, 50, 100, 105, 150, 200, 350, 500, 1000, 5000, 10000, 50000)
        self.print_thresold("Disk inodes", values, "M", 1000000, config.disk_inode)

    def print_thresold(self, name, values, unit, multiplier, configitem):
        header = "%s thresolds (warning=%s, critricial=%s, agressivity=%s):"
        print
        print header % (name, configitem.warning, configitem.critical, configitem.agressivity)
        for val in values:
            val_num = val * multiplier
            warnth = get_thresold_mb(val_num, configitem.warning, configitem.agressivity)
            critth = get_thresold_mb(val_num, configitem.critical, configitem.agressivity)
            print ("%s%s" % (val, unit)).ljust(11), \
                "Warn = %.1f%%" % ((1 - warnth) * 100), \
                ("(%.2f%s Free)" % ((warnth * val), unit)).ljust(8), \
                "Crit = %.1f%%" % ((1 - critth) * 100), \
                ("(%.2f%s Free)" % ((critth * val), unit)).ljust(8)
        

class Commander:
    def __init__(self):
        # To add a new command, add it to this dictionnary :
        # 'short command name':(ClassName to instanciate, 'function to call on ClassName')
        self.commands_mapping = {
            'disk':(DiskCheck, 'check_disk_usage'),
            'disk_inode':(DiskCheck,'check_disk_inode_usage'),
            'mem':(MemoryCheck,'check_memory_usage'),
            'swap':(MemoryCheck,'check_swap_usage'),
            'print':(PrintCheck,'print_thresolds'),
        }
        self.config = Configuration()

    def parse_arguments(self):
        # Construct usage string
        usage = "usage:"
        usage += " %prog "
        for command in self.commands_mapping.keys():
            usage += command + '|'
        # Remove the last '|'
        usage = usage[:-1]

        # Construct OptionParser object
        optparser = OptionParser(usage=usage)
        (options, args) = optparser.parse_args()

        # Validate mendatory arguments
        if len(args) != 1:
            print "Invalid number of arguments"
            optparser.print_help()
            sys.exit(RET_UNKNOWN)
        command = args[0]

        if command not in self.commands_mapping.keys():
            print "Command '%s' unknown." % command
            optparser.print_usage()
            sys.exit(RET_UNKNOWN)
            
        return command

    def engage(self):
        # Call the appropriate function on the appropriate class
        # passing the configuration as a parameter
        command = self.parse_arguments()
        clazz = self.commands_mapping[command][0]
        func = self.commands_mapping[command][1]
        getattr(clazz(), func)(self.config)

if __name__ == '__main__':
    try:
        cmd = Commander()
        cmd.engage()
    except SystemExit, e:
        raise e
    except:
        traceback.print_exc()
        sys.exit(RET_UNKNOWN)

