#!/usr/bin/python3
# Made to act the same as the perl script it's replaced.

import subprocess,json,argparse,re

def main(args):
    try: 
        # Get quota
        cmd = f"ceph -c {args.config} --keyring={args.keyring} osd pool get-quota {args.pool} --format=json"
        res = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE)
        data = json.loads(res.stdout)
    except:
        print(f"Failed to get usage/quota for {args.pool}")
        exit(2)

#    {
#       'current_num_bytes': 1996270682129,
#       'current_num_objects': 507563,
#       'pool_id': 15,
#       'pool_name': 'HCCHC-185946-CLOUD',
#       'quota_max_bytes': 2748779069440,
#       'quota_max_objects': 0
#    }

    if data.get('quota_max_bytes') == 0:
        print("OK: no quota set")
        exit(0)
    if not data.get('current_num_bytes') or not data.get('quota_max_bytes'):
        print("UNKNOWN: No data")
        exit(3)

    warn = parse_human(args.warning)
    crit = parse_human(args.critical)

    bytes_free = data.get('quota_max_bytes') - data.get('current_num_bytes')
    used_perc  = data.get('current_num_bytes') * 100 / data.get('quota_max_bytes')
    free_human = print_human(bytes_free)

    msg = f"{args.pool} {int(used_perc)}% - {free_human} free"

    # values over 100 means bytes, values under 100 means percentages
    if warn > 100 and crit > 100:
        # bytes
        if bytes_free < crit: 
            print(f"CRITICAL: {msg}")
            exit(2)
        if bytes_free < warn:
            print(f"WARNING: {msg}")
            exit(1)
    else:
        # percentages
        if used_perc >= crit:
            print(f"CRITICAL: {msg}")
            exit(2)
        if used_perc >= warn:
            print(f"WARNING: {msg}")
            exit(1)

    # Still here? 
    print(f"OK: {msg}")
    exit(0)


# Python is insane and doesn't have a standard library for converting bytes to/from human input
# "Because it's so simple", so everybody rolls their own untested bodges. This is mine...
# Assumption for both: All 1024 based, none of this 1000-based nonsense :)
# Check is typically called with '100G', without '-B' or '-iB' suffix
def parse_human(size):
    # Convert human to bytes
    units  = {"B": 1, "KB": 2**10, "MB": 2**20, "GB": 2**30, "TB": 2**40}
    size   = size.upper()
    number = re.sub(r'[A-Za-z%\s]', '', size)
    unit   = re.sub(r'[\d\s]+', '', size)
    if unit and unit != "%":
        if not re.search('B$', unit): unit = f"{unit}B" # Add a 'B' to '100G'
        return int(float(number)*units[unit])
    else:
        return int(number)

def print_human(size):
    # Convert bytes to human
    for unit in ("", "K", "M", "G", "T", "P", "E", "Z"):
        if abs(size) < 1024.0:
            return f"{size:3.1f}{unit}"
        size /= 1024.0
    return f"{size:.1f}"

def parse_cmdline():
    parser = argparse.ArgumentParser( description="check_ceph_pool_quota" )
    parser.add_argument('--pool', '-p', type=str, required=True, help="Which pool to check")
    parser.add_argument('--warning', '-w', type=str, default='90', help="Warning treshold")
    parser.add_argument('--critical', '-c', type=str, default='95', help="Critical treshold")
    parser.add_argument('--config', type=str, default='/etc/ceph/ceph.conf', help="Ceph config path")
    parser.add_argument('--keyring', type=str, default='/etc/ceph/ceph.client.admin.keyring', help="Ceph keyring path")
    args = parser.parse_args()
    return args

args = parse_cmdline()
main(args)
