#!/usr/bin/python2
import os, sys
import getopt, datetime
import platform


#automatyczne przeliczenie wartosci od Bajtow w gore
def GetHumanReadable(size,precision=2):
    suffixes=['B','KB','MB','GB','TB']
    suffixIndex = 0
    while size > 1024 and suffixIndex < 3:
        suffixIndex += 1 #increment the index of the suffix
        size = size/1024.0 #apply the division
    return "%.*f%s"%(precision,size,suffixes[suffixIndex])

#sprawdzenie punktu montowania
def getmount(path):        
    path = os.path.realpath(os.path.abspath(path))
    if path == '/boot/efi':
        return path
    while path != os.path.sep:
        if os.path.ismount(path):
            return path
        path = os.path.abspath(os.path.join(path, os.pardir))
    return path

#parametry uzycia dyskow zwracane w formie dictionary
def disk_usage(path):
    st = os.statvfs(path)
    free = st.f_bavail * st.f_frsize
    total = st.f_blocks * st.f_frsize
    used = (st.f_blocks - st.f_bfree) * st.f_frsize
    used_percent=float(used)/total*100
    itotal = st.f_files
    ifree = st.f_ffree
    iused = st.f_files - st.f_ffree
    try: iused_percent=float(iused)/itotal*100
    except:
        iused_percent=1
    return {path: {'total': total,'used': used,'free': free,'used_percent': used_percent, 'itotal': itotal,'ifree': ifree,'iused_percent': iused_percent, 'iused': iused}}

#usuwanie duplikatow w liscie
def make_unique(original_list):
    unique_list = []
    [unique_list.append(obj) for obj in original_list if obj not in unique_list]
    return unique_list


def usage():
    print """python check_disk.py -p [--partition=] -w [--warning=] -c [--critical] -C [--icritical] -W [--iwarning] -m [--megabytes] --gwarning --gcritical
Przyklad:
./check_disk.py -p / -w 10 -c 20 -p /boot -w 11 -c 21 -p /tmp -w 11 -c 22 -p /opt -p /var -p /var/log -p /var/log/audit -W 10 -C 5
Trzeba pamietac zeby progi dawac od razu przy danym fs od lewej do prawej tzn.
./check_disk.py -p / -w 10 -c 20 -p /boot -w 11 -c 21 -p /tmp -W 10 -C 5 --gwarning 10 --gcritical 20 >>>> to jest poprawne
./check_disk.py -p / -w 10 -c 20 -p /boot -p /tmp -W 10 -C 5 --gwarning 10 --gcritical 20 >>> to jest nie poprawne
          """
    sys.exit()

def main():
    partitions_args=[]
    warnings_args=[]
    criticals_args=[]
    mega=0

    try:
        opts, args = getopt.getopt(sys.argv[1:], 'hw:c:p:W:C:m', 
                                   ['help','warning=','critical=','partition=','iwarning=','icritical=','megabytes','gwarning=','gcritical='])
    except getopt.GetoptError:
        usage()
        sys.exit(2)
   
    for opt, arg in opts:
        if opt == '-h':
             usage()
        elif opt in ("-w", "--warning"):
             warnings_args.append(int(arg))
        elif opt in ("-c", "--critical"):
             criticals_args.append(int(arg))
        elif opt in ("-p", "--partition"):
             partitions_args.append(arg)
        elif opt in ("-W", "--iwarning"):
             iwarning=int(arg)
        elif opt in ("-C", "--icritical"):
             icritical=int(arg)
        elif opt in ("-m", '--megabytes'):
             mega=1
        elif opt in ('--gwarning'):
             gwarn=int(arg)
        elif opt in ('--gcritical'):
             gcrit=int(arg)
        else:
             usage()

    part_result=[]
    new_warning=[]
    new_critical=[]
    part_not_distinct=[]    
    for partition in partitions_args:
        part_not_distinct.append(getmount(partition)) # sprawdzenie punktow montowania podanych partycji

 
    part_distinct=make_unique(part_not_distinct) #usuwanie duplikatow punktow montowania
       
    for mountpoint in part_distinct:
        part_result.append(disk_usage(mountpoint)) #sprawdzenie zajetosci per istniejacy punkt montowania
        #print warnings_args[partitions_args.index(mountpoint)]
        try:
           new_warning.append(warnings_args[part_distinct.index(mountpoint)]) # ustalenie progow warning
           new_critical.append(criticals_args[part_distinct.index(mountpoint)]) # ustalenie progow critical
        except IndexError:
           new_warning.append(gwarn) # ustalenie progow global warning
           new_critical.append(gcrit) # ustalenie progow global critical
    
    perfdata=""
    outputOK=""
    outputWARN=""
    outputCRIT=""
    outputINODE=""
    i=0
    crit=0   
    warn=0
 
    try: gwarn,gcrit
    except NameError as e:
        pass
    
    #wyswietlanie i wyliczanie
    if mega == 0: #procent used
       for element in part_result:
           for tag,value in element.items():
               p_used=value['used_percent']
               
               if p_used > float(new_critical[i]):
                  outputCRIT+=tag+"="+GetHumanReadable(value['free'])+" "
                  crit=1
               elif p_used > float(new_warning[i]):
                  outputWARN+=tag+"="+GetHumanReadable(value['free'])+" "
                  warn=1
               else:
                  outputOK+=tag+"="+GetHumanReadable(value['free'])+" "             
                #sprawdzenie inode used 
               if value['iused_percent'] > float(icritical):
                  outputINODE+=tag+" InodeCRIT "+format(value['iused_percent'],'.2f')+" "
                  crit=1
               elif value['iused_percent'] > float(iwarning):
                  outputINODE+=tag+" InodeWARN "+format(value['iused_percent'],'.2f')+" "
                  warn=1

               warning=float(new_warning[i])/100*value['total']/1024
               critical=float(new_critical[i])/100*value['total']/1024
               perfdata+=tag+"="+str(value['used']/1024)+"KB;"+format(warning,'.0f')+";"+format(critical,'.0f')+";0;"+str(value['total']/1024)+"; "
               #output+=tag+"="+GetHumanReadable(value['used'])+" "
               i+=1
    elif mega == 1: #megabajty used
       for element in part_result:
           for tag,value in element.items():
               used=value['used']/1024/1024
               if used < new_critical[i]:
                  outputCRIT+=tag+"="+GetHumanReadable(value['free'])+" "
                  crit=1
               elif used < new_warning[i]:
                  outputWARN+=tag+"="+GetHumanReadable(value['free'])+" "
                  warn=1
               else:
                  outputOK+=tag+"="+GetHumanReadable(value['free'])+" "
                 
                #sprawdzenie inode used
               if value['iused_percent'] > float(icritical):
                  outputINODE+=tag+" InodeCRIT "+format(value['iused_percent'],'.2f')+" "
                  crit=1
               elif value['iused_percent'] > float(iwarning):
                  outputINODE+=tag+" InodeWARN "+format(value['iused_percent'],'.2f')+" "
                  warn=1

               perfdata+=tag+"="+str(value['used']/1024)+"KB;"+str(new_warning[i]*1024)+";"+str(new_critical[i]*1024)+";0;"+str(value['total']/1024)+"; "
               #output+=tag+"="+GetHumanReadable(value['used'])+" "
               i+=1

    if crit==1:
       print "DISK CRITICAL Free Space "+outputCRIT+" "+outputINODE+"| "+perfdata
       sys.exit(2)
    elif warn==1:
       print "DISK WARNING Free Space "+outputWARN+" "+outputINODE+"| "+perfdata
       sys.exit(1)
    else:
       print "DISK OK Free Space "+outputOK+"| "+perfdata
       sys.exit(0)

if __name__ == '__main__':
    main()
