#!/usr/bin/env python
#
# 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 itertools import chain


# A note about MAX_SOCKETS:
#
# According to the MIB, MAX_SOCKETS should be 64. However, since CheckMK 1.6
# doesn't appear to have something like 2.*'s SNMPTree functionality, we
# brute-force by explicitly querying every possible socket. Unfortunately,
# this means querying ~3K OIDs. In practice, the data received from the customer
# shows that they only have two or four sockets per infeed, so I've set the
# MAX_SOCKETS here to 16, which cuts the queried OIDs to 800. Using MAX_SOCKETS
# of 4 would still satisfy the customer's current needs, but leaves no leeway
# if the get more sockets in an infeed.
MAX_TOWERS  = 4
MAX_INFEEDS = 4
MAX_SOCKETS = 16 # XXX

TOWER_ARR_SIZE  = MAX_TOWERS
INFEED_ARR_SIZE = MAX_TOWERS * MAX_INFEEDS
SOCKET_ARR_SIZE = MAX_TOWERS * MAX_INFEEDS * MAX_SOCKETS

TOWER_NUM_OFFSET   = 0
INFEED_NUM_OFFSET  = TOWER_NUM_OFFSET   + 1
INFEED_ID_OFFSET   = INFEED_NUM_OFFSET  + TOWER_ARR_SIZE
INFEED_NAME_OFFSET = INFEED_ID_OFFSET   + INFEED_ARR_SIZE
INFEED_VOLT_OFFSET = INFEED_NAME_OFFSET + INFEED_ARR_SIZE
SOCKET_NUM_OFFSET  = INFEED_VOLT_OFFSET + INFEED_ARR_SIZE
SOCKET_ID_OFFSET   = SOCKET_NUM_OFFSET  + INFEED_ARR_SIZE
SOCKET_NAME_OFFSET = SOCKET_ID_OFFSET   + SOCKET_ARR_SIZE
SOCKET_LOAD_OFFSET = SOCKET_NAME_OFFSET + SOCKET_ARR_SIZE


def parse_sentry_pdu(info):
    outlets_info = {}
    results = info[0]
    num_towers = int(results[TOWER_NUM_OFFSET])

    for tower_id in range(num_towers):
        num_infeeds = int(results[INFEED_NUM_OFFSET + tower_id])

        for infeed_id in range(num_infeeds):
            infeed_idx = tower_id * MAX_TOWERS + infeed_id

            infeed_sid  = results[INFEED_ID_OFFSET + infeed_idx]
            infeed_name = results[INFEED_NAME_OFFSET + infeed_idx]
            infeed_voltage = float(results[INFEED_VOLT_OFFSET + infeed_idx]) / 10
            num_outlets = int(results[SOCKET_NUM_OFFSET + infeed_idx])

            for outlet_id in range(num_outlets):
                outlet_idx = MAX_INFEEDS * MAX_SOCKETS * tower_id + MAX_SOCKETS * infeed_id + outlet_id
                outlet_sid  = results[SOCKET_ID_OFFSET + outlet_idx]
                outlet_name = results[SOCKET_NAME_OFFSET + outlet_idx]
                outlet_load = float(results[SOCKET_LOAD_OFFSET + outlet_idx]) / 100

                outlet_key = '%s.%s.%s' % (tower_id, infeed_id, outlet_id)
                outlets_info[outlet_key] = {
                    'infeed_id':   infeed_sid,
                    'infeed_name': infeed_name,
                    'infeed_voltage': infeed_voltage,
                    'outlet_id':   outlet_sid,
                    'outlet_name': outlet_name,
                    'outlet_load': outlet_load,
                }

    return outlets_info 


# 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 (3, 'item not found in snmp data')

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

    if voltage < 0:
        return (1, 'Infeed voltage unavailable')

    if params['type'] == 'infeed':
        return check_metric(params, 'volts', voltage)

    elif params['type'] == 'outlet':
        if amps < 0:
            return (1, 'Outlet load unavailable')
        return check_metric(params, 'watts', voltage * amps)


def check_metric(params, metric_name, metric_value):
    crit_metric_above = params.get('crit_%s_above' % metric_name)
    warn_metric_above = params.get('warn_%s_above' % metric_name)
    warn_metric_below = params.get('warn_%s_below' % metric_name)
    crit_metric_below = params.get('crit_%s_below' % metric_name)
        
    state = 0

    if crit_metric_above and crit_metric_above < metric_value:
        state = 2
    elif crit_metric_below and crit_metric_below > metric_value:
        state = 2
    elif warn_metric_above and warn_metric_above < metric_value:
        state = 1
    elif warn_metric_below and warn_metric_below > metric_value:
        state = 1

    return (state, '%.1f %s' % (metric_value, metric_name))


# Inventory function, returning inventory based upon SNMP parsed result above
def inventory_sentry_pdu(parsed):
    items = []

    for id, outlet in parsed.items():
        plug_name = '%s %s [infeed %s] power' % (
            outlet['outlet_id'],
            outlet['outlet_name'],
            outlet['infeed_id']
        )
        infeed_name = 'Infeed %s %s' % (
            outlet['infeed_id'],
            outlet['infeed_name']
        )

        items.append((plug_name,   { 'id': id, 'type': 'outlet' }))
        items.append((infeed_name, { 'id': id, 'type': 'infeed' }))

    return items

 
check_info['sentry_pdu_outlets_power'] = {
    'parse_function': parse_sentry_pdu,
    'check_function': check_sentry_pdu,
    'inventory_function': inventory_sentry_pdu,
    'service_description': 'Outlet %s',
    'group': 'sentry_pdu_outlets_power',
    'snmp_info': (
        '.1.3.6.1.4.1.1718.3',
        # In CheckMK 2.* there is SMPTree, but 1.6 doesn't seem to have that.
        # Therefore we resort to this sledge-hammer approach to get the
        # information we need. It's... not ideal.
        list(chain(
            # Number of towers
            ['1.4.0'],
            
            # Number of infeeds
            # .1.3.6.1.4.1.1718.3.2.1.1.5.<tower #>
            ['2.1.1.5.%s' % (x) for x in range(1, MAX_TOWERS+1)],

            # Infeed IDs:
            # .1.3.6.1.4.1.1718.3.2.2.1.2.<tower #>.<infeed #>
            ['2.2.1.2.%s.%s'  % (x, y) for x in range(1, MAX_TOWERS+1)
                                       for y in range(1, MAX_INFEEDS+1)],

            # Infeed names:
            # .1.3.6.1.4.1.1718.3.2.2.1.3.<tower #>.<infeed #>
            ['2.2.1.3.%s.%s'  % (x, y) for x in range(1, MAX_TOWERS+1)
                                       for y in range(1, MAX_INFEEDS+1)],

            # Infeed voltage:
            # .1.3.6.1.4.1.1718.3.2.2.1.11.<tower #>.<infeed #>
            ['2.2.1.11.%s.%s' % (x, y) for x in range(1, MAX_TOWERS+1)
                                       for y in range(1, MAX_INFEEDS+1)],

            # Number of outlets
            # .1.3.6.1.4.1.1718.3.2.2.1.9.<tower #>.<infeed #>
            ['2.2.1.9.%s.%s'  % (x, y) for x in range(1, MAX_TOWERS+1)
                                       for y in range(1, MAX_INFEEDS+1)],

            # Outlet IDs:
            # .1.3.6.1.4.1.1718.3.2.3.1.2.<tower #>.<infeed #>.<outlet #>
            ['2.3.1.2.%s.%s.%s' % (x, y, z) for x in range(1, MAX_TOWERS+1)
                                            for y in range(1, MAX_INFEEDS+1)
                                            for z in range(1, MAX_SOCKETS+1)],

            # Outlet names:
            # .1.3.6.1.4.1.1718.3.2.3.1.3.<tower #>.<infeed #>.<outlet #>
            ['2.3.1.3.%s.%s.%s' % (x, y, z) for x in range(1, MAX_TOWERS+1)
                                            for y in range(1, MAX_INFEEDS+1)
                                            for z in range(1, MAX_SOCKETS+1)],

            # Outlet load:
            # .1.3.6.1.4.1.1718.3.2.3.1.7.<tower #>.<infeed #>.<outlet #>
            ['2.3.1.7.%s.%s.%s' % (x, y, z) for x in range(1, MAX_TOWERS+1)
                                            for y in range(1, MAX_INFEEDS+1)
                                            for z in range(1, MAX_SOCKETS+1)],
        ))),
    'snmp_scan_function': lambda oid: 'Sentry Switched -48 VDC' in oid('.1.3.6.1.2.1.1.1.0')
}
