#!/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_cpu.py
#  Author - Ian Perry <iperry@indigex.com>
#
#  Purpose:  Tell what the script does
#
#  Version History:
#       2022.11.04 - Initial Creation
#       2022.11.07 - Modified to use lowercase u instead of U
#       2025.01.15 - Modified to correct from average to current
# ***************************************************************************
try:
    from inmon_utils import *
except:
    import sys
    print("Failed to import inmon_utils")
    sys.exit(3)

import argparse
import time # Necessary for SSH communication
import os

import paramiko

parser = argparse.ArgumentParser(description="Checks CPU on ")

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",
    required=True
)
parser.add_argument(
    "-w", "--warning",
    dest="warning",
    help="Specify warning threshold",
    required=True
)

args = parser.parse_args()

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

# Create SSH client via paramiko, set host keys and auto add policy
client = paramiko.client.SSHClient()
client.load_system_host_keys()
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

# Connect to remote host.
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']
    }
)

# First run to get /proc/stat info. /proc/stat lists since last boot
stdin, stdout, stderr = client.exec_command('cat /proc/stat | grep \'^cpu \'') # nosec: properly sanitized
before_output = stdout.readlines() # readlines() reads until the process terminates
if stdout.channel.recv_exit_status() != 0: # Handle non-zero exit status
    client.close()
    err = stderr.read().decode()
    print(f"[UNKNOWN] - Error while retrieving /proc/stat data: {err}")
    sys.exit(3)
stdin.close() # Close the stdin. Prevents a paramiko issue with closing client before it's finished

time.sleep(3) # Wait some period of time before gathering again. Below block does the same thing as above.

stdin, stdout, stderr = client.exec_command('cat /proc/stat | grep \'^cpu \'') # nosec: properly sanitized
after_output = stdout.readlines()
if stdout.channel.recv_exit_status() != 0:
    client.close()
    err = stderr.read().decode()
    print(f"[UNKNOWN] - Error while retrieving /proc/stat data: {err}")
    sys.exit(3)
stdin.close()

client.close() # Close the client.

# Split the line into subsequent strings and discard 'cpu'
before_output = before_output[0].split()[1:]
after_output = after_output[0].split()[1:]

# Cast all to ints
before_output = [int(i) for i in before_output]
after_output = [int(i) for i in after_output]

# Subtract before from after to get the values for 3 second spread
output = [after_output[i] - before_output[i] for i in range(len(after_output))]

# CPU total
cpu_total = sum(output)

# CPU idle
cpu_idle = output[3]

# Divide idle by total to get fraction
frac = cpu_idle/cpu_total

# Subtract previous fraction from 1.0 to get time spent not being idle
frac = 1.0 - frac

# Multiply by 100 to get a percentage
res = round(float(frac * 100),2)

if res < float(args.warning):
    retval = ["OK", 0]
elif res >= float(args.warning) and res < float(args.critical):
    retval = ["WARNING", 1]
elif res > float(args.critical):
    retval = ["CRITICAL", 2]
else:
    retval = ["UNKNOWN", 3]

print(f"[{retval[0]}] - CPU usage is {res}% | CPUusage={res}%;{args.warning};{args.critical}")
sys.exit(retval[1])