#!/usr/bin/python3

import os
import sys
import psutil
import re
import shutil

def kmgt(sz, frac=1):
    t = {
        'K': pow(1024, 1),
        'M': pow(1024, 2),
        'G': pow(1024, 3),
        'T': pow(1024, 4),
        '': 1}

    for k in sorted(t, key=t.__getitem__, reverse=True):
        if sz >= t[k]:
            n = sz / float(t[k])
            tpl = "{:." + str(frac) + "f}{}"
            return tpl.format(n, k)


if len(sys.argv) == 1:
    method = 'info'
else:
    method = sys.argv[1]

if method == 'info':
    info = {
        'Protocol': '0.1',
        'Description': 'Free disk space',
        'Version': '0.1',
        'Methods': 'check info makeconfig'
    }
    for k, v in info.items():
        print("{}: {}".format(k, v))

elif method == 'check':
    # read parameters
    prefix = os.getenv('PREFIX')
    maxlim = int(os.getenv('MAXLIM','90'))
    ignore_device = os.getenv('IGNORE_DEVICE','^/dev/loop')
    ignore_mount = os.getenv('IGNORE_MOUNT','')  # '^/media/|^/snap/'

    for p in psutil.disk_partitions():
        # print("?", p)
        if ignore_device:
            if re.search(ignore_device, p.device):
                #print("ignore device", ignore_device)
                continue
        if ignore_mount:
            if re.search(ignore_mount, p.mountpoint):
                #print("ignore device", ignore_device)
                continue
        u = shutil.disk_usage(p.mountpoint)
        perc = u.used * 100 / u.total
        total = kmgt(u.total)
        used = kmgt(u.used)
        free = kmgt(u.free)

        data="""
NAME: {prefix}df-{mount}
TAGS: la
METHOD: numerical|maxlim={maxlim}
DETAILS: {perc:.2f}%, {used}/{total} used, {free} free
STATUS: {perc:.2f}
""".lstrip().format(prefix=prefix, maxlim=maxlim, mount=p.mountpoint, perc=perc, used=used, total=total, free=free)
        print(data)

elif method=='makeconfig':
    data="""
# Alert if any disk used for more then this %%
MAXLIM=90

# Ignore partitions where devices matches regex 
IGNORE_DEVICE=^/dev/loop

# Ignore partitions where devices matches regex 
IGNORE_MOUNT=^/media/|^/tmp/

### COMMON SETTINGS

# PREFIX2=df:

# Policy if not default
POLICY=
""".strip()

    print(data)