#!/usr/bin/python -W ignore::DeprecationWarning
# Michele Baldessari - Leitner Technologies - 2011
# 06.09.2011

import datetime
import getopt
import os
import pprint
import subprocess
import sys
import traceback
from pysnmp.entity.rfc3413.oneliner import cmdgen
from pysnmp.proto.api import v2c

SNMP_FNHAOID = "1.3.6.1.4.1.12356.1.100"
SNMP_FNHA_SERIAL = SNMP_FNHAOID + ".6.1.2"
SNMP_FNHA_NAME = SNMP_FNHAOID + ".6.1.11"
SNMP_PORT = 161

RET_OK = 0
RET_WARNING = 1
RET_CRITICAL = 2
RET_UNKNOWN = 3

def snmp_dict_string(d):
    s = ""
    for i in d.keys():
        s = s + "[%s] %s" % (i, d[i])
    return s

def check_fortigate_ha(host, community, basedir):
    # if entry for cluster does not exist yet create it and return success
    # if it does exist check the file and verify the order is correct
    oid = v2c.ObjectIdentifier(SNMP_FNHAOID)
    auth = cmdgen.CommunityData('foobar', community)
    errorIndication, errorStatus, errorIndex, \
                     varBindTable = cmdgen.CommandGenerator().nextCmd(auth,
                         cmdgen.UdpTransportTarget((host, SNMP_PORT)),
                         (oid))

    if errorIndication:
        print "SNMP Query error"
        raise SystemExit, RET_UNKNOWN

    if errorStatus:
        print '%s at %s\n' % ( errorStatus.prettyPrint(),
            errorIndex and varBindTable[-1][int(errorIndex)-1] or '?')
        raise SystemExit, RET_UNKNOWN


    forti_ha_dict = {}
    for varBindTableRow in varBindTable:
        for name, val in varBindTableRow:
            forti_ha_dict[name.prettyPrint()] = val

    current_cluster = []
    try:
        master = {}
        master[forti_ha_dict[SNMP_FNHA_SERIAL + ".1"]] = forti_ha_dict[SNMP_FNHA_NAME + ".1"]
        current_cluster.append(master)
    except:
        print "Unable to retrieve %s HA oids" % (SNMP_FNHA_SERIAL + ".1")
        raise SystemExit, RET_UNKNOWN

    try:
        slave = {}
        slave[forti_ha_dict[SNMP_FNHA_SERIAL + ".2"]] = forti_ha_dict[SNMP_FNHA_NAME + ".2"]
        current_cluster.append(slave)
    except:
        print "Only Master %s is active" % (snmp_dict_string(master))
        raise SystemExit, RET_CRITICAL
        
    fname = os.path.join(basedir, host)
    if not os.path.exists(fname): # First time we look at the cluster (we assume that the master is the correct one) and we exit 0 without messages
        f = open(fname, "w+")
        for i in current_cluster:
            for j in i.keys():
                f.write("%s:%s\n" % (j, i[j]))
        f.close()
    else: # We already stored the serials of the cluster
        f = open(fname, "r")
        lines = f.readlines()
        if len(lines) > 2:
            print "More than two cluster members in file found: %s" % (len(lines))
            raise SystemExit, RET_UNKNOWN

        cycle = [0, 1]
        for i in cycle:
            for current_serial in current_cluster[i].keys():
                if current_serial != lines[i].strip().split(":")[0]:
                    print "WARNING - Cluster has changed priorities"
                    raise SystemExit, RET_WARNING
                if current_cluster[i][current_serial] != lines[i].strip().split(":")[1]:
                    print "WARNING - Cluster has changed priorities"
                    raise SystemExit, RET_WARNING
    
        f.close()
        print "OK - Cluster is in healthy state"
        raise SystemExit, RET_OK


def usage():
    print ("fortinet_ha.py -h|--help -b|--basedir=<base directory> -C|--community=<community> -H|--host=<host>")

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hb:H:C:", ["help", "basedir=", "host=", "community="])
    except getopt.GetoptError as err:
        usage()
        sys.exit(2)

    community = 'public'
    host = 'localhost'
    basedir = "/etc/nagios3/fortiserial/"
    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
            sys.exit()
        elif o in ("-b", "--basedir"):
            basedir = a
        elif o in ("-H", "--host"):
            host = a
        elif o in ("-C", "--community"):
            community = a
        else:
            assert False, "unhandled option"

    if host != None:
        check_fortigate_ha(host, community, basedir)
    else:
        usage()
        sys.exit(2)  


if __name__ == "__main__":
    main()