import subprocess

class SimpleSNMP():
    '''Minimal SNMP call from python'''

    def walk(self, options):
        command = "snmpwalk {0}".format(options)
        walk = subprocess.Popen(command.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        output, err = walk.communicate()
        return(output)


class FortiGet:
    '''Simple class to fetch snmp infos from fortigate'''

    def __init__(self):
        self.lastfield = lambda output: output.split()[-1]


    def get_data(self, snmpclient, user, password, host, oid):
        walk = snmpclient.walk("-v3 -u {0} -a MD5 -A {1} -l authPriv -X {2} {3} {4}".format(user, password, password, host, oid))
        return(walk)


    def get_all_vpn_ids(self, snmpclient, user, password, host):
        walk = self.get_data(snmpclient, user, password, host, "FORTINET-FORTIGATE-MIB::fgVpnTunEntIndex")
        ids = [x.split()[-1] for x in walk.split("\n") if x != ""]
        return(ids)


    def get_quantity_vdom(self, snmpclient, user, password, host):
        quantity = self.get_data(snmpclient, user, password, host, "FORTINET-FORTIGATE-MIB::fgVdNumber")
        return(self.lastfield(quantity))


    def get_max_vdoms(self, snmpclient, user, password, host):
        maxvdoms = self.get_data(snmpclient, user, password, host, "FORTINET-FORTIGATE-MIB::fgVdMaxVdoms")
        return(self.lastfield(maxvdoms))


    def get_vpn_status(self, vpnid, snmpclient, user, password, host):
        oid = "FORTINET-FORTIGATE-MIB::fgVpnTunEntStatus.{0}".format(vpnid) 
        status = self.get_data(snmpclient, user, password, host, oid)
        return(self.lastfield(status))


    def get_vpn_phase_name(self, vpnid, phase, snmpclient, user, password, host):
        oid = "FORTINET-FORTIGATE-MIB::fgVpnTunEntPhase{0}Name.{1}".format(phase, vpnid)
        vpnname = self.get_data(snmpclient, user, password, host, oid)
        return(self.lastfield(vpnname))


    def get_vdomid(self, vpnid, snmpclient, user, password, host):
        oid = "FORTINET-FORTIGATE-MIB::fgVpnTunEntVdom.{0}".format(vpnid)
        vdom = self.get_data(snmpclient, user, password, host, oid)
        return(self.lastfield(vdom))


    def get_vdom_name(self, vdomid, snmpclient, user, password, host):
        oid = "FORTINET-FORTIGATE-MIB::fgVdEntName.{0}".format(vdomid)
        vdomname = self.get_data(snmpclient, user, password, host, oid)
        return(self.lastfield(vdomname))