#!/usr/bin/python3
import psutil
import os
import argparse
from pprint import pprint

headers = 'Mountpoint','Used','Total','Usage'

suffixes = ['B', 'KB', 'MB', 'GB', 'TB', 'PB']
def humansize(nbytes):
    i = 0
    while nbytes >= 1024 and i < len(suffixes)-1:
        nbytes /= 1024.
        i += 1
    f = ('%.2f' % nbytes).rstrip('0').rstrip('.')
    return '%s %s' % (f, suffixes[i])


def parserDefinintion():
    parser = argparse.ArgumentParser(
                        prog = 'check_disk.py',
                        description = 'What the program does',
                        epilog = 'Text at the bottom of help')

    parser.add_argument('-w', '--warning',default=80,type=int)
    parser.add_argument('-c', '--critical',default=90,type=int)
    parser.add_argument('-x', '--excludefs')

    return parser.parse_args()   

def checkDisk(args):

    hdd = psutil.disk_usage('/')
    hdd = psutil.disk_partitions(1)

    # set arguments
    warning = args.warning
    critical = args.critical

    if args.excludefs != None:
        skip_types = args.excludefs.split(',')
    else:
        skip_types = []

    output = []

    for entry in hdd:

        if entry.fstype in skip_types:
            continue

        stats  = psutil.disk_usage(entry.mountpoint)

        if stats.total == 0:
            continue
        
        if(stats.percent > critical):
            print("check_disk CRITICAL: %s %s%% bytes usage (%s/%s)" % (entry.mountpoint,stats.percent,humansize(stats.used),humansize(stats.total)))
            exit(2)

        if(stats.percent > warning):
            print("check_disk WARNING: %s %s%% bytes usage (%s/%s)" % (entry.mountpoint,stats.percent,humansize(stats.used),humansize(stats.total)))
            exit(1)

        inodes = os.statvfs(entry.mountpoint)
        inodes_used = (1 - (inodes.f_ffree / inodes.f_files)) * 100

        if(inodes_used > warning):
            print("check_disk WARNING: %s %.2f%% inodes usage (%s/%s)" % (entry.mountpoint,inodes_used,inodes.f_files-inodes.f_ffree,inodes.f_files))
            exit(1)

        if(inodes_used > critical):
            print("check_disk CRITICAL: %s %.2f%% inodes usage (%s/%s)" % (entry.mountpoint,inodes_used,inodes.f_files-inodes.f_ffree,inodes.f_files))
            exit(2)

    print("check_disk OK: Disk usage and inode usage under %.f%%" % (warning))
    exit(0)


if __name__ == "__main__":

    args = parserDefinintion()

    checkDisk(args)