#!/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_ubiquiti_uap.py
#  Author - Ian Perry <iperry@indigex.com>
#
#  Purpose:  Gathers information from Unifi UAPs
#
#  Version History:
#       2023.04.12 - 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
import re

parser = argparse.ArgumentParser(description="Checks statistics on Unifi UAPs")

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(
    "-a", "--icinga",
    dest="icinga",
    help="Specify icinga device hostname if different from FQDN"
)
parser.add_argument(
    "-u", "--username",
    dest="username",
    help="Username to use for SSH",
    required=True
)
parser.add_argument(
    "-d", "--disk",
    dest="disk",
    help="Thresholds for disk usage (warn:crit)",
    default="10:10"
)
parser.add_argument(
    "-c", "--cpu",
    dest="cpu",
    help="Thresholds for cpu usage (warn:crit)",
    default="80:90"
)
parser.add_argument(
    "-r", "--ram",
    dest="ram",
    help="Thresholds for ram usage (warn:crit)",
    default="90:95"
)
parser.add_argument(
    "-p", "--parent",
    dest="parent",
    help="FQDN of device passive checks will be submitted to.",
    required=True
)

# Parse arguments
args = parser.parse_args()

# Ping hostname and immediately exit if we can't reach.
if not ping(args.hostname):
    print(f"CRITICAL - Unable to reach {args.hostname}")
    sys.exit(2)

# Expand user in keyfile
keyfile = os.path.expanduser(args.keyfile)

# Split into individual variables.
disk_warning, disk_critical = args.disk.split(":")
cpu_warning, cpu_critical = args.cpu.split(":")
ram_warning, ram_critical = args.ram.split(":")

# Prepare array for passive check results
passive_results = []

# PREPARE THRESHOLDS
# HANDLE RAM
regex_percent = re.compile('^[0-9]+$')
regex_absolute = re.compile('^[0-9]+[BKMGbkmg]$')


if regex_percent.match(ram_warning) and regex_percent.match(ram_critical):
    mode = "PCT"
elif regex_absolute.match(ram_warning) and regex_absolute.match(ram_critical):
    mode = "ABS"
    warning_denomination = ram_warning[-1]
    critical_denomination = ram_critical[-1]

    warning_amount = float(ram_warning[0:-1])
    critical_amount = float(ram_critical[0:-1])

    if re.match(r'b|B', warning_denomination):
        warning_level = warning_amount*(1024**0)
    elif re.match(r'k|K', warning_denomination):
        warning_level = warning_amount*(1024**1)
    elif re.match(r'm|M', warning_denomination):
        warning_level = warning_amount*(1024**2)
    elif re.match(r'g|G', warning_denomination):
        warning_level = warning_amount*(1024**3)
    else:
        passive_results += PassiveCheckResult(
            args.parent, args.hostname, "RAM Usage", 3,
            f"[UNKNOWN] - Bad unit of measurement used for absolute memory warn level. Warn level {ram_warning}, crit level {ram_critical}"
        )

    if re.match(r'b|B', critical_denomination):
        critical_level = critical_amount*(1024**0)
    elif re.match(r'k|K', critical_denomination):
        critical_level = critical_amount*(1024**1)
    elif re.match(r'm|M', critical_denomination):
        critical_level = critical_amount*(1024**2)
    elif re.match(r'g|G', critical_denomination):
        critical_level = critical_amount*(1024**3)
    else:
        passive_results += PassiveCheckResult(
            args.parent, args.hostname, "RAM Usage", 3,
            f"[UNKNOWN] - Bad unit of measurement used for absolute memory crit level. Warn level {ram_warning}, crit level {ram_critical}"
        )
else:
    passive_results += PassiveCheckResult(
        args.parent, args.hostname, "RAM Usage", 3,
        f"[UNKNOWN] - Arguments must be same type, either both percent or both absolute."
    )


# HANDLE CPU
if cpu_warning > cpu_critical:
    print("UNKNWON - Warning level must be less than or equal to critical level.")

# MAIN

# Create SSH client and prepare it.
client = paramiko.client.SSHClient()
client.load_system_host_keys()
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

# Connect to device
client.connect(hostname, username=args.username, key_filename=keyfile, disabled_algorithms={
               'keys': ['rsa-sha2-256', 'rsa-sha2-512'], 'pubkeys': ['rsa-sha2-256', 'rsa-sha2-512']})

# Run and get output for CPU
stdin, stdout, stderr = client.exec_command('cat /proc/stat | grep \'^cpu \'')
exit_status = stdout.channel.recv_exit_status()
cpu_output = stdout.read().decode()

# Run and get output for memory
stdin, stdout, stderr = client.exec_command('cat /proc/meminfo')
exit_status = stdout.channel.recv_exit_status()
cpu_output = stdout.read().spllitlines()

# Run and get output for disk usage
stdin, stdout, stderr = client.exec_command('df /tmp')
exit_status = stdout.channel.recv_exit_status()
cpu_output = stdout.read().splitlines

# Close the client.
client.close()