#!/usr/bin/env python3
#-*- encoding: utf-8
#
# Author : Alexander Vogel (alexander.vogel.2305@gmail.com)
# Date   : 2026-10-19
# License: GNU General Public License v2
#
# Libexec: ExtremeCloud IQ Controller


import argparse
import json
import requests
import sys
import urllib3


urllib3.disable_warnings()


def parse_arguments(argv):
    parser = argparse.ArgumentParser(description=__doc__)
    
    parser.add_argument("--server", type=str, required=True)
    parser.add_argument("--port", type=int, default=5825)
    parser.add_argument("--username", type=str, required=True)
    parser.add_argument("--password", type=str, required=True)
    parser.add_argument("--sections", type=str, required=False, default="all")

    args = parser.parse_args(argv)
    
    return args

 
class ExtremeXiqControllerAgent():
    def __init__(self, server: str, port: int, username: str, password: str) -> None:
        
        self.server = server
        self.port = port
        self.username = username
        self.password = password
        
        self.url = f"https://{args.server}:{args.port}"
        self.token = ""
        
        # try to get access token
        self.token = self._authenticate()


    def _authenticate(self):

        payload = {
            "grantType": "password",
            "userId": self.username,
            "password": self.password
        }
        
        resp = self._request("POST", f"/management/v1/oauth2/token", data=payload)
        
        try:
            return resp.get("access_token")
        except:
            raise Exception(f"API-Error {resp.status_code}: {resp.text}")


    def _request(self, method: str, endpoint, headers=None, data={}):

        api_url = self.url + endpoint
        
        if not headers:
            headers = {
                'Accept': 'application/json',
                'Authorization': f'Bearer {self.token}',
                'Content-Type': 'application/json'
            }

        resp = requests.request(method, api_url, headers=headers, data=json.dumps(data), verify=False)

        if not resp.ok:
            raise RuntimeError(f"API-Error {resp.status_code}: {resp.text}")

        return resp.json()


    def get_access_points(self):
        return self._request("GET", f"/management/v1/aps/query")
        

    def get_access_point_state(self):
        return self._request("GET", f"/management/v1/state/aps")


    def get_all_services(self):
        return self._request("GET", f"/management/v1/services")


    def get_sites(self):
        return self._request("GET", f"/management/v3/sites")


    def get_all_stations(self):
        return self._request("GET", f"/management/v2/stations/query")


if __name__ == "__main__":

    args = parse_arguments(sys.argv[1:])

    # init api
    agent_netextreme_xiq_controller = ExtremeXiqControllerAgent(
        server = args.server,
        port = args.port,
        username = args.username,
        password = args.password
    )

    # APs (Access Points)
    if "ap" in args.sections or "all" in args.sections:
        access_points = agent_netextreme_xiq_controller.get_access_points()
        states = agent_netextreme_xiq_controller.get_access_point_state()
        
        if len(access_points) > 0:
            sys.stdout.write(f"<<<netextreme_xiq_controller_ap:sep(0)>>>\n")
        
        for ap in access_points:

            parsed = {
                "name": ap['apName'],
                "power": ap['pwrUsage'],
                "uptime": ap['sysUptime'],
                "clients": 0,
                "ipAddress": ap['ipAddress'],
                "hardwareType": ap['hardwareType'],
                "serialNumber": ap['serialNumber'],
                "softwareVersion": ap['softwareVersion'],
                "macAddress": ap['macAddress'],
                "description": ap['description'],
            }

            for r in ap['radios']:
                parsed['clients'] += r['clients']

            # state
            for state in states:
                if state['apSerialNo'] == ap['serialNumber']:   
                    parsed['status'] = state['entityStatus']['operationalStatus']
                    break
            
            sys.stdout.write(f"{parsed}\n")

    # determine all clients when WLAN or site data is required. This data can then be used to determine the number of 
    # clients per WLAN or site. This is faster than querying each site or WLAN for clients via individual API requests.
    if 'site' in args.sections or 'wlan' in args.sections or 'all' in args.sections:
        stations = agent_netextreme_xiq_controller.get_all_stations()

    # Sites (Locations)
    if 'site' in args.sections or 'all' in args.sections:
        sites = agent_netextreme_xiq_controller.get_sites()
        
        if len(sites) > 0:
            sys.stdout.write("<<<netextreme_xiq_controller_site:sep(0)>>>\n")

            for s in sites:

                parsed = {
                    "name": s['siteName'],
                    "clients": 0,
                }

                # determine clients
                for st in stations:
                    if st['siteName'] == s['siteName']:
                        parsed['clients'] +=1

                sys.stdout.write(f"{parsed}\n")

    # WLANs
    if 'wlan' in args.sections or 'all' in args.sections:
        wlans = agent_netextreme_xiq_controller.get_all_services()

        if len(wlans) > 0:
            sys.stdout.write("<<<netextreme_xiq_controller_wlan:sep(0)>>>\n")
        
            for wlan in wlans:

                parsed = {
                    "serviceName": wlan['serviceName'],
                    "status": wlan['status'],
                    "ssid": wlan['ssid'],
                    "clients": 0,
                }
                
                # determine clients
                for st in stations:
                    if st['serviceName'] == wlan['serviceName']:
                        parsed['clients'] +=1    
                
                sys.stdout.write(f"{parsed}\n")