#!/usr/bin/env python3
# ***************************************************************************
#  IIIIII NNN  NNN  Copyright (C) 2022 Innovative Networks, Inc.
#    II    NNN  N   All Rights Reserved. Any redistribution or reproduction
#    II    N NN N   of part or all of the content of this program in any form
#    II    N  NNN   without expressed written consent of the copyright holder
#  IIIIII NNN  NNN  is strictly prohibited.  Please contact admins@in-kc.com
#   Be Innovative.  for additional information.
# ***************************************************************************
#  check_uap_disk.py
#  Author - Ian Perry <iperry@indigex.com>
#
#  Purpose:  Checks disk usage on ubiquiti access points
#
#  Version History:
#       2022.11.07 - Initial Creation
# ***************************************************************************
try:
    from inmon_utils import *
except:
    import sys
    print("Failed to import inmon_utils")
    sys.exit(3)

import argparse
import paramiko
import time # Necessary for SSH communication
import os

parser = argparse.ArgumentParser(description="Checks disk on Ubiquiti access points")

parser.add_argument(
    "-i", "--keyfile",
    dest="keyfile",
    help="Specify SSH keyfile",
    required=True
)
parser.add_argument(
    "-H", "--hostname",
    dest="hostname",
    help="Specify device hostname",
    required=True
)
parser.add_argument(
    "-u", "--username",
    dest="username",
    help="Specify username to log into device with",
    required=True
)
parser.add_argument(
    "-c", "--critical",
    dest="critical",
    help="Specify critical threshold",
    default=10
)
parser.add_argument(
    "-w", "--warning",
    dest="warning",
    help="Specify warning threshold",
    default=10
)

args = parser.parse_args()

keyfile = os.path.expanduser(args.keyfile)

client = paramiko.client.SSHClient()
client.load_system_host_keys()
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
client.connect(args.hostname, username=args.username, key_filename=keyfile, disabled_algorithms={'keys': ['rsa-sha2-256', 'rsa-sha2-512'], 'pubkeys': ['rsa-sha2-256', 'rsa-sha2-512']})
stdin, stdout, stderr = client.exec_command('df /tmp')
time.sleep(1)
client.close()
output = stdout.read().splitlines()
output = [i.decode() for i in output]
first = output[0].split()
second = output[1].split()
out_dict = {}
for i in range(len(second)):
    out_dict[first[i]] = second[i]

disk_total = int(out_dict["Available"]) + int(out_dict["Used"])
percent_free = round((int(out_dict["Available"])*100)/disk_total, 2)
percent_used = round((int(out_dict["Used"])*100)/disk_total, 2)

if percent_free > int(args.warning):
    status = "[OK]"
    code = 0
elif percent_free <= int(args.warning) and percent_free > int(args.critical):
    status = "[WARNING]"
    code = 1
elif percent_free <= int(args.critical):
    status = "[CRITICAL]"
    code = 2
else:
    status = "[UNKNOWN]"
    code = 3

status += f" - Disk usage: {out_dict['Used']}K/{out_dict['Available']}K; {percent_free}% available. | disk_used={percent_used}%"
print(status)
sys.exit(code)