#!/usr/bin/env python3
# ***************************************************************************
#  IIIIII NNN  NNN  Copyright (C) 2023 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.
# ***************************************************************************
#  inmon_utils.py
#  Author - Ian Perry <iperry@indigex.com>
#
#  Purpose:  Utility functionality for Python INmon scripts
#
#  Version History:
#       2023.04.26 - Modified to add INmon passive check API functionality.
# ***************************************************************************
try:
    from inmon_utils import *
except:
    import sys
    print("Failed to import inmon_utils")
    sys.exit(3)
import argparse
import paramiko
import logging
import sys
import re
import xml.etree.ElementTree as ET

parser = argparse.ArgumentParser(description="Obtains zenarmor status for an OPNsense firewall")

parser.add_argument(
    "-v",
    "--verbose",
    dest="verbose",
    action="store_true",
    help="sets logging level to info"
)

parser.add_argument(
    "-i",
    "--keyfile",
    dest="keyfile",
    help="ssh keyfile",
    required=True
)

parser.add_argument(
    "-H",
    "--hostname",
    dest="hostname",
    help="hostname of the device to check",
    required=True
)

parser.add_argument(
    "-U",
    "--username",
    dest="username",
    help="username to log into the device under",
    required=True
)

args = parser.parse_args()

def is_running(text):
    # Search for " <1-5 numbers>."
    m = re.findall(r'\ (\d{1,5})\.', text)
    if len(m) == 0:
        return((False, False))
    else:
        return((True, m[0]))

client = paramiko.client.SSHClient()
client.load_system_host_keys()
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
client.connect(str(args.hostname), username=str(args.username), key_filename=str(args.keyfile))

stdin, stdout, stderr = client.exec_command ('/bin/cat /conf/config.xml')
exit_status = stdout.channel.recv_exit_status()
if exit_status != 0:
    print(f"Error: {stderr.read().decode()}")
    sys.exit(3)
root = ET.fromstring(stdout.read().decode())

exit_status = 0
exit_text = ""

database = root.find('./OPNsense/Zenarmor/general/database')
if database is None:
    print("Error: Zenarmor configuration not found")
    sys.exit(3)
database_type = database.find('Type').text
if bool(database.find('Remote').text):
    exit_text += "Zenarmor DB running on a remote host. "
elif database_type == "MN":
    stdin, stdout, stderr = client.exec_command('/usr/local/bin/sudo /usr/sbin/service mongod status')
    exit_status_mongod = stdout.channel.recv_exit_status()
    mongod_outlines = stdout.readline().strip()
    database_status, database_pid = is_running(mongod_outlines)
    if database_status:
        exit_text += f"Mongo running with pid {database_pid}. "
    else:
        exit_status = max(exit_status, 2)
        exit_text += "Mongo is not running. "
elif database_type == "ES":
    exit_status = max(exit_status, 3)
    exit_text += "Check has not been configured for local Elasticsearch. Please contact iperry for implementation. "    

stdin, stdout, stderr = client.exec_command('/usr/local/bin/sudo /usr/sbin/service eastpect status')
exit_status_zenarmor = stdout.channel.recv_exit_status()
if exit_status_zenarmor != 0:
    print(f"Error: {stderr.read().decode()}")
    sys.exit(3)
zenarmor_outlines = stdout.readline().strip()
zenarmor = is_running(zenarmor_outlines)

client.close()

if not zenarmor[0]:
    exit_status = max(exit_status, 2)
    exit_text += "Zenarmor is not running. "
else:
    exit_text += f"Zenarmor is running with pid {zenarmor[1]}."

header = "[CRITICAL] - " if exit_status == 2 else "[OK] - "

print(f"{header}{exit_text}")
sys.exit(exit_status)