#!/usr/bin/env python3
#
# Copyright 2024 Spearhead Systems SRL
#
# Docs for this system:
# https://cdn10.servertech.com/assets/documents/documents/135/original/manual_CDU_Y-30932L.pdf
# https://cdn10.servertech.com/assets/documents/documents/793/original/Sentry3.mib

from cmk.base.plugins.agent_based.agent_based_api.v1 import (
    register,
    Service,
    Result,
    State,
    SNMPTree,
    contains,
    OIDEnd,
)


def extract_outlet(tower_id, infeed_id, outlet_id, data):
    return next(filter(lambda x: x[0] == f'{tower_id+1}.{infeed_id+1}.{outlet_id+1}', data))[1]


# SNMP parsing function
def parse_sentry_pdu(string_table):
    outlets_info = {}
    num_towers = int(string_table[0][0][0])

    for tower_id in range(num_towers):
        num_infeeds = int(string_table[1][0][tower_id])

        for infeed_id in range(num_infeeds):
            infeed_sid = string_table[2][infeed_id][tower_id]
            infeed_voltage = float(string_table[3][infeed_id][tower_id]) / 10
            num_outlets = int(string_table[4][infeed_id][tower_id])

            for outlet_id in range(num_outlets):
                outlet_sid  = extract_outlet(tower_id, infeed_id, outlet_id, string_table[5])
                outlet_name = extract_outlet(tower_id, infeed_id, outlet_id, string_table[6])
                outlet_load = float(extract_outlet(tower_id, infeed_id, outlet_id, string_table[7])) / 100

                outlets_info[f'{tower_id}.{infeed_id}.{outlet_id}'] = {
                  'infeed_id': infeed_sid,
                  'infeed_voltage': infeed_voltage,
                  'outlet_id': outlet_sid,
                  'outlet_name': outlet_name,
                  'outlet_load': outlet_load
                }

    return outlets_info


# Inventory function, returning inventory based upon SNMP parsed result above
def discovery_sentry_pdu(section):
    for key, outlet in section.items():
        name = f'{outlet["outlet_id"]} {outlet["outlet_name"]} (infeed {outlet["infeed_id"]}) power'
        yield Service(item=name, parameters={ 'id': key })


# Check function, returning warn/crit based upon SNMP parsed result above
def check_sentry_pdu(item, params, section):
    outlet = section.get(params['id'])
    if outlet is None:
        return

    voltage = outlet['infeed_voltage']
    amps = outlet['outlet_load']

    if voltage < 0:
        yield Result(state=State.WARN, summary='Infeed voltage unavailable')
        return
    elif amps < 0:
        yield Result(state=State.WARN, summary='Outlet load unavailable')
        return
        
    power = voltage * amps
    state = State.OK

    crit_watts_above = params.get('crit_watts_above')
    warn_watts_above = params.get('warn_watts_above')
    warn_watts_below = params.get('warn_watts_below')
    crit_watts_below = params.get('crit_watts_below')

    if crit_watts_above and crit_watts_above < power:
        state = State.CRIT
    elif crit_watts_below and crit_watts_below > power:
        state = State.CRIT
    elif warn_watts_above and warn_watts_above < power:
        state = State.WARN
    elif warn_watts_below and warn_watts_below > power:
        state = State.WARN

    yield Result(state=state, summary=f'{power:.1f} watts')


register.snmp_section(
    name='sentry_pdu_outlets_power',
    parse_function=parse_sentry_pdu,
    fetch=[
       # Number of towers (integer from 0 to 4)
       SNMPTree(
           base='.1.3.6.1.4.1.1718.3.1.4',
           oids=['0']
       ),

       # Number of infeeds (integer from 0 to 4):
       # .1.3.6.1.4.1.1718.3.2.1.1.5.<tower #>
       SNMPTree(
           base='.1.3.6.1.4.1.1718.3.2.1.1.5',
           oids=['1', '2', '3', '4'],
       ),

       # Infeed IDs:
       # .1.3.6.1.4.1.1718.3.2.2.1.2.<tower #>.<infeed #>
       SNMPTree(
           base='.1.3.6.1.4.1.1718.3.2.2.1.2',
           oids=['1', '2', '3', '4'],
       ),

       # Infeed voltage:
       # .1.3.6.1.4.1.1718.3.2.2.1.11.<tower #>.<infeed #>
       SNMPTree(
           base='.1.3.6.1.4.1.1718.3.2.2.1.11',
           oids=['1', '2', '3', '4'],
       ),

       # Number of outlets (integer from 0 to 64):
       # .1.3.6.1.4.1.1718.3.2.2.1.9.<tower #>.<infeed #>
       SNMPTree(
           base='.1.3.6.1.4.1.1718.3.2.2.1.9',
           oids=['1', '2', '3', '4'],
       ),

       # Outlet IDs:
       # .1.3.6.1.4.1.1718.3.2.3.1.2.<tower #>.<infeed #>.<outlet #>
       SNMPTree(
           base='.1.3.6.1.4.1.1718.3.2.3.1',
           oids=[OIDEnd(), '2'],
       ),

       # Outlet names:
       # .1.3.6.1.4.1.1718.3.2.3.1.3.<tower #>.<infeed #>.<outlet #>
       SNMPTree(
           base='.1.3.6.1.4.1.1718.3.2.3.1',
           oids=[OIDEnd(), '3'],
       ),

       # Outlet load:
       # .1.3.6.1.4.1.1718.3.2.3.1.7.<tower #>.<infeed #>.<outlet #>
       SNMPTree(
           base='.1.3.6.1.4.1.1718.3.2.3.1',
           oids=[OIDEnd(), '7']
        ),
    ],
    detect=contains('.1.3.6.1.2.1.1.1.0', 'Sentry Switched -48 VDC'),
)

register.check_plugin(
    name='sentry_pdu_outlets_power',
    service_name='Outlet %s',
    discovery_function=discovery_sentry_pdu,
    check_function=check_sentry_pdu,
    check_default_parameters={},
    check_ruleset_name='sentry_pdu_outlets_power',
)