#!/usr/bin/env python
#
print("""
################################################################################
# The goal of this program is check F5 disk space. It's a Linux installation
# with multiple partitions in a nutshell. If any partition is filled up to 90%,
# an email alarm will be sent out.
#
################################################################################
# Version: 0.1, Date: August 14 2019
# Author: Sam Li, yangsec888@gmail.com, x.x.1972
# Usage:
#         $ dsk-chk --node xxx
###############################################################################
# Prerequisite:
# a. F5 remote SSH access service are setup.
# b. You have privileged access to the F5 file system.
###############################################################################
""")
#from client import F5Client
import f5_admin
from f5_admin.util import *
import argparse
import os
import time
from datetime import datetime
from os.path import isfile, join
# send notice over relay host
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.header import Header

# command argument switch setup
parser = argparse.ArgumentParser(description='This program will monitor F5 disk space utilization for you. ')
parser.add_argument('-n','--node', help='(Optional) Remote F5 hostname or IP address',required=False)
parser.add_argument('-a','--all',help='(Optional) Exam all known F5 nodes in the local cache directory. ', action='store_true', default=False, required=False)
parser.add_argument('-u','--user', help='(Optional) F5 box logon user ID. Default to root. ',required=False)
parser.add_argument('-p','--password', nargs='+', help='(Optional) F5 box root password',required=False)
parser.add_argument('-e','--email', help='Email notiification address when alarm trip',required=True)
parser.add_argument('-r','--refresh',help='(Optional) Refresh the F5 running configuration in the local cache directory. ', action='store_true', default=False, required=False)
parser.add_argument('-v','--verbose', help='(Optional) Verbose mode',action='store_true', default=False, required=False)
args = parser.parse_args()


# disk space check worker
def chk_worker(f5_node):
    partition_tally=0
    with f5_admin.F5Client(credential, 7, args.verbose) as client:
        current_time = int(time.time())
        client.load(f5_node)
        my_conn = client.ssh_connect()      # setup the connection
        if client.credential["user_name"] == "root":    # root user by default have bash shell once logon
            my_cmd = "df -k"
        else:
            my_cmd = "bash -c 'df -k'"                  # non-root user by default have tmsh once logon
        result = client.ssh_command(my_conn,my_cmd, "")
        partition_list = [i for i in result if "%" in i and not "Filesystem" in i]
        partition_tally = len(partition_list)
        print("Total number of partition checked: ", partition_tally, "\n")
        #  Partition disk space examination
        for x in partition_list:
            if "/var/apm/mount" in x:
                continue
            dsk_ent = x.split()
            print('Checking on partition: ', dsk_ent)
            if '%' in dsk_ent[-2]:
                u = int(dsk_ent[-2].split('%')[0])
                if u > 90: # 90% threshhold
                    print("")
                    if args.verbose:
                        print("Partition usage details: ", x)
                    notice = "F5 Node: " + client.node  + " disk partition filled: " + x
                    print("Notice: ", notice)
                    send_notice(args.email, notice)
    return partition_tally

# send out email notice
def send_notice(to_email,notice):
    print("Sending notice to: ", to_email)
    # build the email content
    msg = MIMEMultipart('alternative')
    msg['Subject'] = "F5 Disk Partition Utilization Warning @wmstools01"
    msg['From'] = str(Header('Deploy Engineer <deploy@apps.mask.com>'))
    msg['To'] = to_email
    mime_part = MIMEText(notice, 'plain')
    msg.attach(mime_part)
    # send out email
    server=smtplib.SMTP("exrelay.us.mask.com:25")
    server.sendmail(msg['From'],msg['To'], msg.as_string())
    server.quit()
    print("Email sent. ")

################################################################################
#      Main Program
################################################################################
if __name__ == '__main__':

    # simple command argument check
    if not (args.node or args.all):
        print("Command for help: ",__file__, " -h")
        exit(1)
    # Setup F5 connection string
    if args.password:
        credential = set_credential_1(args.user, *args.password)
    else:
        credential = set_credential_2(args.user)

    total_chk = 0
    if args.node:
        sub_total = chk_worker(args.node)
        total_chk += sub_total
    elif args.all:
        with f5_admin.F5Client(credential, 7, args.verbose) as client:
            nodes = client.get_node_list()
            whitelist = ['xx-rhwebdev1', 'xx-rhwebdev2', 'xx-bigiq1', 'xx-toolsprd1', 'xx-toolsprd2']
            chk_nodes = [ node for node in nodes if node not in whitelist]
        for node in chk_nodes:
            sub_total = chk_worker(node)
            total_chk += sub_total

    # Print the confirmation
    print("Total number of disk partitions checked: ", total_chk)
    print("""
All done. Bye!""")
