#!/usr/bin/python

import os, sys, time, random, glob, traceback

opt_verbose = '-v' in sys.argv
opt_force   = '-f' in sys.argv

omd_site = os.environ['OMD_SITE']
omd_root = os.environ['OMD_ROOT']

config_file       = omd_root + '/etc/diskspace.conf'
plugin_dir        = omd_root + '/share/diskspace'
plugin_dir_local  = omd_root + '/local/share/diskspace'

# Initial configuration
min_free_bytes = None
max_file_age   = None
min_file_age   = None

plugins = {}

def error(s):
    sys.stderr.write('ERROR: %s\n' % s)

def terminate(s):
    error(s)
    sys.exit(1)

def log(s):
    sys.stdout.write('%s\n' % s)

def verbose(s):
    if opt_verbose:
        log(s)

def read_config():
    try:
        execfile(config_file, globals(), globals())
    except IOError:
        pass # ignore non existant config
    except Exception, e:
        terminate('Invalid configuration: %s' % e)

def resolve_paths():
    for plugin in plugins.values():
        resolved = []
        for path in plugin.get('cleanup_paths', []):
            # Make relative paths absolute ones
            if path[0] != '/':
                path = omd_root + '/' + path

            # This resolves given path patterns to really existing files.
            # It also ensures that the files in the resolved list do really exist.
            resolved += glob.glob(path)

        if resolved:
            plugin['cleanup_paths'] = resolved
        elif 'cleanup_paths' in plugin:
            del plugin['cleanup_paths']

def load_plugins():
    try:
        local_plugins = os.listdir(plugin_dir_local)
    except OSError:
        local_plugins = [] # this is optional

    plugin_files = [ p for p in os.listdir(plugin_dir) if p not in local_plugins ]

    for base_dir, file_list in [ (plugin_dir, plugin_files), (plugin_dir_local, local_plugins) ]:
        for f in file_list:
            if f[0] == '.':
                continue

            plugins[f] = {}

            path = base_dir + '/' + f
            verbose('Loading plugin: %s' % path)
            try:
                execfile(path, plugins[f], plugins[f])
            except Exception, e:
                error('Exception while loading plugin "%s": %s' % (path, e))

    # Now transform all path patterns to absolute paths for really existing files
    resolve_paths()

def collect_file_infos():
    for plugin in plugins.values():
        for path in plugin.get('cleanup_paths', []):
            result = os.stat(path)
            plugin.setdefault('file_infos', {})[path] = (result.st_size, result.st_mtime)

def fmt_bytes(b):
    b = float(b)
    base = 1024
    if b >= base * base * base * base:
        return '%.2fTB' % (b / base / base / base / base)
    elif b >= base * base * base:
        return '%.2fGB' % (b / base / base / base)
    elif b >= base * base:
        return '%.2fMB' % (b / base / base)
    elif b >= base:
        return '%.2fkB' % (b / base)
    else:
        return '%.0fB' % b

def get_free_space():
    # FIXME: Take eventual root reserved space into account
    for l in os.popen('df -P -B1 ' + omd_root).readlines():
        if l[0] == '/':
            vol, size_bytes, used_bytes, free_bytes, used_perc, mp = l.split()
            return int(free_bytes)

def above_threshold(b):
    return b >= min_free_bytes

def delete_file(path, reason):
    try:
        log('Deleting file (%s): %s' % (reason, path))
        os.unlink(path)
        return True
    except Exception, e:
        error('Error while deleting %s: %s' % (path, e))
    return False

# Loop all files to check wether or not files are older than
# max_age. Simply remove all of them.
def cleanup_aged():
    if max_file_age is None:
        verbose('Not cleaning up too old files (max_file_age not configured)')
        return
    max_age = time.time() - max_file_age

    for plugin in plugins.values():
        for path, (size, mtime) in plugin.get('file_infos', {}).items():
            if mtime < max_age:
                if delete_file(path, 'too old'):
                    del plugin['file_infos'][path]
            else:
                verbose('Not deleting %s' % path)

def oldest_candidate(file_infos):
    if file_infos:
        # Sort by modification time
        sorted_infos = sorted(file_infos.items(), key = lambda i: i[1][1])
        oldest = sorted_infos[0]
        if oldest[1][1] < time.time() - min_file_age:
            return oldest[0]

def cleanup():
    if min_file_age is None:
        terminate('Not cleaning up oldest files of plugins (min_file_age not configured). terminating.')

    # the scheduling of the cleanup job is supposed to be equal for
    # all sites. To ensure that not only one single site is always
    # cleaning up, we add a a random wait before cleanup.
    sleep_sec = float(random.randint(0, 10000)) / 1000
    verbose('Sleeping for %0.3f seconds' % sleep_sec)
    time.sleep(sleep_sec)

    # Loop all cleanup plugins to find the oldest candidate per plugin
    # which is older than min_age and delete this file.
    for plugin_name, plugin in plugins.items():
        oldest = oldest_candidate(plugin.get('file_infos', {}))
        if oldest:
            delete_file(oldest, plugin_name + ': my oldest')

def main():
    load_plugins()
    collect_file_infos()

    # get used diskspace of the sites volume
    bytes_free = get_free_space()
    verbose('Free space: %s' % fmt_bytes(bytes_free))

    cleanup_aged()

    bytes_free = get_free_space()
    verbose('Free space (after max_age cleanup): %s' % fmt_bytes(bytes_free))

    # check diskspace against configuration
    if not opt_force and above_threshold(bytes_free):
        # -> ok: exit
        verbose('Free space is above threshold of %s. Nothing to be done.' % fmt_bytes(min_free_bytes))
        return

    # free diskspace is below threshold, start cleanup
    cleanup()

# #############################################################################

read_config()

if min_free_bytes is None:
    verbose('minimal free bytes (min_free_bytes) not configured. terminating.')
    sys.exit(0)

try:
    main()
except SystemExit:
    raise
except:
    terminate('Unexpected exception: %s' % traceback.format_exc())
