#!/usr/bin/env python3

import subprocess
import argparse
import sys

class CheckAP:
    def __init__(self, names_oid, status_oid, version, community, controller):
        self._names_oid = names_oid
        self._status_oid = status_oid
        self._snmp_version = version
        self._snmp_community = community
        self._controller = controller

        self._snmp_call = [
            "snmpwalk",
            "-On",
            "-v", self._snmp_version,
            "-c", self._snmp_community,
            self._controller,
        ]
        self._ap_oids = self.get_oids()

    def get_oids(self):
        """
        Get OIDs of the available access points
        Return name -> oid dict
        """

        output = subprocess.check_output(self._snmp_call + [self._names_oid])

        ret = dict()
        for line in output.decode().split("\n"):
            line = line.strip()
            if not line:
                continue

            oid, _, _, name = line.split(" ")
            name = name.strip('"')
            oid_part = oid.partition(self._names_oid)[2]

            ret[name] = oid_part

        return ret

    def get_ap_status(self, ap_name):
        """
        Query an AP's status by its name
        Return True or False
        """
        oid = self._status_oid + self._ap_oids[ap_name]
        output = subprocess.check_output(self._snmp_call + [oid])
        status = int(output.decode().split(" ")[3])

        return status == 1


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-n", "--name", help="AP name", required=True)
    parser.add_argument("-c", "--community", help="SNMP Community", required=True)
    parser.add_argument("-C", "--controller", help="Controller IP Address", required=True)
    parser.add_argument("-v", "--version", help="SNMP Version", default="2c")
    args = parser.parse_args()

    check = CheckAP(
        community = args.community,
        controller = args.controller,
        version = args.version,
        names_oid = "1.3.6.1.4.1.14823.2.2.1.5.2.1.4.1.3",
        status_oid = "1.3.6.1.4.1.14823.2.2.1.5.2.1.4.1.19",
    )

    if check.get_ap_status(args.name):
        print("{} OK - associated with controller".format(args.name))
        sys.exit(0)
    else:
        print("{} CRITICAL - not associated with controller ".format(args.name))
        sys.exit(2)
