#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2017, Florent Thiery
'''
Criticality: Normal
Checks that partitions are in conformity with expected norms, and that sufficient free space is available
'''
import subprocess
import sys
import os

YELLOW = '\033[93m'
GREEN = '\033[92m'
RED = '\033[91m'
DEF = '\033[0m'


PATHS = [
    {
        'mount_point': '/',
        'recommended_types': ('ext4', 'zfs'),
        'min_size_gbytes': 9,
        'reco_size_gbytes': 14,
        'min_available_gbytes': 4,
    },
    {
        'mount_point': '/home/msuser/msinstance',
        'recommended_types': ('ext4', 'zfs', 'nfs', 'nfs4'),
        'min_size_gbytes': 5,
        'reco_size_gbytes': 300,
        'min_available_gbytes': 5,
    },
    {
        'mount_point': '/home/skyreach',
        'recommended_types': ('ext4', 'zfs', 'nfs', 'nfs4'),
        'min_size_gbytes': 5,
        'reco_size_gbytes': 9,
        'min_available_gbytes': 2,
    },
    {
        'type': 'memory',
        'min_size_gbytes': 2,
        'reco_size_gbytes': 4,
    },
]


def to_gbytes(size_bytes):
    return int(round(size_bytes / (1024 * 1024 * 1024)))


def get_memory_gbytes():
    memory_gbytes = 0
    with open('/proc/meminfo', 'r') as f:
        for l in f:
            if 'MemTotal:' in l:
                memory = l.split('MemTotal:')[1].strip()
                memory_kbytes, unit = memory.split(' ')
                if unit != 'kB':
                    print('Warning, unexpected unit %s.' % unit)
                memory_gbytes = int(round(int(memory_kbytes) / (1024 * 1024)))
    if not memory_gbytes:
        print('Failed to get memory size.')
    return memory_gbytes


def get_path(path):
    # Filesystem     Type   1B-blocks        Used   Available
    # /dev/loop2     ext4 52710469632 38253940736 11755397120
    status, output = subprocess.getstatusoutput('df --output="source,fstype,size,avail" -B 1 %s | tail -n 1' % path)
    if status == 0:
        dev, fstype, size, available = output.split()
    else:
        dev = fstype = size = available = None
    return dev, fstype, to_gbytes(int(size)), to_gbytes(int(available))


def check_allocation(dev):
    root_dev = os.path.basename(dev)[:3]
    if not root_dev:
        return True
    dev_partitions = list()
    with open('/proc/partitions', 'r') as f:
        for l in f:
            if root_dev in l:
                dev_partitions.append(l)

    max_size = 0
    total_size = 0
    for p in dev_partitions:
        major, minor, blocks, name = p.split()
        size = int(blocks) * 512
        if name == root_dev:
            max_size = size
            if root_dev.startswith('md'):
                total_size += size
        else:
            total_size += size
    unallocated = max_size - total_size
    unallocated_gbytes = to_gbytes(unallocated)
    if unallocated_gbytes > 1:
        print('Warning: %s%s GB are unallocated on %s.%s' % (YELLOW, unallocated_gbytes, root_dev, DEF))
        return False
    return True


if __name__ == '__main__':
    error = False
    warning = False
    for p in PATHS:
        psize = None
        if p.get('mount_point'):
            mount_point = p['mount_point']
            if os.path.exists(mount_point):
                mount_point = os.path.realpath(mount_point)
                name = 'Partition of %s' % mount_point
                dev, fstype, psize, available = get_path(mount_point)
                if fstype not in p.get('recommended_types'):
                    print('Warning: %s fs type not recommended %s(current: %s, recommended: %s)%s.' % (name, YELLOW, fstype, p['recommended_types'], DEF))
                    warning = True
                if 'nfs' not in fstype:
                    warning = not check_allocation(dev)
                min_available_gbytes = p.get('min_available_gbytes')
                if min_available_gbytes and available < min_available_gbytes:
                    print('%s has less than %s GB available %s(%s GB available)%s.' % (name, min_available_gbytes, RED, available, DEF))
                    error = True
                else:
                    print('%s has %s GB available.' % (name, available))
            else:
                print('%s not found, cannot check.' % mount_point)
        elif p.get('type') == 'memory':
            name = 'Memory'
            psize = get_memory_gbytes()

        if psize:
            if psize < p['min_size_gbytes']:
                print('%s is smaller than the minimum required size %s(%s GB < %s GB)%s.' % (name, RED, psize, p['min_size_gbytes'], DEF))
                error = True
            elif psize < p['reco_size_gbytes']:
                print('%s is smaller than the recommended size %s(%s GB < %s GB)%s.' % (name, YELLOW, psize, p['reco_size_gbytes'], DEF))
                warning = True
            else:
                print('%s is bigger than recommended size %s(%s GB >= %s GB)%s.' % (name, GREEN, psize, p['reco_size_gbytes'], DEF))

    if error:
        print('Errors found.')
        code = 1
    elif warning:
        print('Some warnings were found.')
        code = 3
    else:
        print(GREEN + 'All OK.' + DEF)
        code = 0

    sys.exit(code)