from pysnmp import hlapi
from influxdb import InfluxDBClient


def get(target, oids, credentials, port=161, engine=hlapi.SnmpEngine(), context=hlapi.ContextData()):
    handler = hlapi.getCmd(
        engine,
        credentials,
        hlapi.UdpTransportTarget((target, port)),
        context,
        *construct_object_types(oids)
    )
    return fetch(handler, 1)[0]


def construct_object_types(list_of_oids):
    object_types = []
    for oid in list_of_oids:
        object_types.append(hlapi.ObjectType(hlapi.ObjectIdentity(oid)))
    return object_types


def fetch(handler, count):
    result = []
    for i in range(count):
        try:
            error_indication, error_status, error_index, var_binds = next(handler)
            if not error_indication and not error_status:
                items = {}
                for var_bind in var_binds:
                    items[str(var_bind[0])] = cast(var_bind[1])
                result.append(items)
            else:
                raise RuntimeError('Got SNMP error: {0}'.format(error_indication))
        except StopIteration:
            break
    return result


def cast(value):
    try:
        return int(value)
    except (ValueError, TypeError):
        try:
            return float(value)
        except (ValueError, TypeError):
            try:
                return str(value)
            except (ValueError, TypeError):
                pass
    return value


KVW_TOTAL = '1.3.6.1.4.1.31034.12.1.1.2.6.1.1.2.0.1'  # SCHLEIFENBAUER-DATABUS-MIB::sdbDevInKWhTotal.0.1
KVW_SUBTOTAL = '1.3.6.1.4.1.31034.12.1.1.2.6.1.1.3.0.1'  # SCHLEIFENBAUER-DATABUS-MIB::sdbDevInKWhSubtotal.0.1
POWER_FACTOR = '1.3.6.1.4.1.31034.12.1.1.2.6.1.1.4.0.1'  # SCHLEIFENBAUER-DATABUS-MIB::sdbDevInPowerFactor.0.1
ACTUAL_CURRENT = '1.3.6.1.4.1.31034.12.1.1.2.6.1.1.5.0.1'  # SCHLEIFENBAUER-DATABUS-MIB::sdbDevInActualCurrent.0.1
PEAK_CURRENT = '1.3.6.1.4.1.31034.12.1.1.2.6.1.1.6.0.1'  # SCHLEIFENBAUER-DATABUS-MIB::sdbDevInPeakCurrent.0.1
ACTUAL_VOLTAGE = '1.3.6.1.4.1.31034.12.1.1.2.6.1.1.7.0.1'  # SCHLEIFENBAUER-DATABUS-MIB::sdbDevInActualVoltage.0.1
MIN_VOLTAGE = '1.3.6.1.4.1.31034.12.1.1.2.6.1.1.8.0.1'  # SCHLEIFENBAUER-DATABUS-MIB::sdbDevInMinVoltage.0.1
VOLT_AMPERE = '1.3.6.1.4.1.31034.12.1.1.2.6.1.1.9.0.1'  # SCHLEIFENBAUER-DATABUS-MIB::sdbDevInPowerVoltAmpere.0.1
POWER_WATT = '1.3.6.1.4.1.31034.12.1.1.2.6.1.1.10.0.1'  # SCHLEIFENBAUER-DATABUS-MIB::sdbDevInPowerWatt.0.1

hosts = [
    {'host': 'pdu01', 'ip': '10.100.199.1'},
    {'host': 'pdu02', 'ip': '10.100.199.2'},
    {'host': 'pdu03', 'ip': '10.100.199.3'}
]

oids = [
    KVW_TOTAL,
    KVW_SUBTOTAL,
    POWER_FACTOR,
    ACTUAL_CURRENT,
    PEAK_CURRENT,
    ACTUAL_VOLTAGE,
    MIN_VOLTAGE,
    VOLT_AMPERE,
    POWER_WATT,
]
community = '5up3r5ecr3t'

host = 'localhost'
port = 8086
user = 'root'
password = 'root'
dbname = 'mydb'

client = InfluxDBClient(host, port, user, password, dbname)

while True:

    for host in hosts:
        try:
            data = get(host['ip'], oids, hlapi.CommunityData(community))

            json_body = [
                {
                    'measurement': 'pdu',
                    'tags': {
                        'host': host.get('host'),
                    },
                    'fields': {
                        "kwh_total" : data.get(KVW_TOTAL),
                        "kwh_subtotal": data.get(KVW_SUBTOTAL),
                        "power_factor": int(data.get(POWER_FACTOR)) / 100,
                        "peak_current": int(data.get(PEAK_CURRENT)) / 100,
    sleep(10)