#!/usr/bin/env python3
#
# Send an SNMP trap to a destination.
#

import pscheduler
import pysnmp
from pysnmp.hlapi import *
import re
import sys

# check for missing required fields
def missing_input(data_json, test_type, level):

    try:
        # general fields
        dest = data_json['dest']
        trap_oid = data_json['trap-oid']
    except KeyError:
        return True

    try:
        community = data_json['_community']
    except KeyError:
        try:
            security_name = data_json['security-name']
            security_level = data_json['security-level'].lower()

            # authentication and encryption
            if security_level == 'authpriv':
                auth_protocol = data_json['auth-protocol']
                auth_key = data_json['_auth-key']
                priv_protocol = data_json['priv-protocol']
                priv_key = data_json['_priv-key']

            # encryption
            if security_level == 'authnopriv':
                auth_protocol = data_json['auth-protocol']
                auth_key = data_json['_auth-key']
        except KeyError:
            return True

    except KeyError:
        return True

    return False


def build_tuple(oid):

    if re.match(r'^((\.\d)|\d)+(\.\d+)*$', oid) is None:
        try:
            temp = oid.split('::') # change this
            args = [temp[0]]
            temp = temp[1].split('.')
            args.extend(temp)
            try:
                obj_id = (args[0], args[1], args[2])
            except IndexError:
                obj_id = (args[0], args[1])
        except IndexError:
            pscheduler.fail('Incomplete/Invalid OID')

    else:
        obj_id = (oid,)

    return obj_id

def get_credentials(data_json):

    try:
        return CommunityData(data_json['_community'])
    except KeyError:
        auth_protocols = { 'sha': usmHMACSHAAuthProtocol,
                           'md5': usmHMACMD5AuthProtocol
                         }

        priv_protocols = { 'des': usmDESPrivProtocol,
                           '3des': usm3DESEDEPrivProtocol,
                           'aes': usmAesCfb128Protocol,
                           'aes128': usmAesCfb128Protocol,
                           'aes192': usmAesCfb192Protocol,
                           'aes256': usmAesCfb256Protocol
                         }

        if data_json['security-level'].lower() == 'noauthnopriv':
            return UsmUserData(data_json['security-name'])

        elif data_json['security-level'].lower() == 'authpriv':
            return UsmUserData(data_json['security-name'], data_json['_auth-key'], data_json['_priv-key'],
                                authProtocol=auth_protocols[data_json['auth-protocol'].lower()],
                                privProtocol=priv_protocols[data_json['priv-protocol'].lower()])

        elif data_json['security-level'].lower() == 'authnopriv':
            return UsmUserData(data_json['security-name'], data_json['_auth-key'],
                              authProtocol=auth_protocols[data_json['auth-protocol'].lower()])

        else:
            pscheduler.fail("Failed to archive: unrecognized security level: %s" % data_json['security-level'])


def get_generator(data):

    trap_object = build_tuple(data['trap-oid'])

    notification = NotificationType(ObjectIdentity(*trap_object))
    credentials = get_credentials(data)

    try:
        instance_index = data['instance-index']
        notification._NotificationType__instanceIndex += (instance_index,)
    except KeyError:
        pass

    try: 
        varbinds = data['trap-varbinds']
        for varbind in varbinds:
            obj = build_tuple(varbind['oid'])
            notification._NotificationType__objects[obj] = varbind['value']
    except KeyError:
        pass

    return sendNotification(SnmpEngine(),
                            credentials,
                            UdpTransportTarget((data['dest'], 162)),
                            ContextData(),
                            'trap',
                            notification
    )

def archive(input_json):

    try:
        g = get_generator(json['data'], json['result']['data'], json['result']['oid'])
        errorIndication, errorStatus, errorIndex, varBinds = next(g)
        if errorIndication:
            succeeded = False
            data = None
            error = "Failed to archive: %s" % str(errorIndication).strip('\n')
        else:
            succeeded = True
            data = None
            error = None

    except (pysnmp.error.PySnmpError, pysnmp.smi.error.NoSuchObjectError) as e:
        succeeded = False
        data = None
        error = "Failed to archive: %s" % str(e).strip('\n')

    output_json = {'succeeded': succeeded,
                    'data': data,
                    'error': error}


def archive(json):

    try:
        g = get_generator(json['data'])
        errorIndication, errorStatus, errorIndex, varBinds = next(g)
        if errorIndication:
            succeeded = False
            data = None
            error = "Failed to archive: %s" % str(errorIndication).strip('\n')
        else:
            succeeded = True
            data = None
            error = None

    except (pysnmp.error.PySnmpError, pysnmp.smi.error.NoSuchObjectError) as e:
        succeeded = False
        data = None
        error = "Failed to archive: %s" % str(e).strip('\n')

    return {
        'succeeded': succeeded,
        'data': data,
        'error': error
    }



PARSER = pscheduler.RFC7464Parser(sys.stdin)
EMITTER = pscheduler.RFC7464Emitter(sys.stdout)

for parsed in PARSER:
    try:
        EMITTER(archive(parsed))
    except BrokenPipeError as ex:
        log.warning("Broken pipe during archiving; parent must have exited.")
        pscheduler.succeed()

pscheduler.succeed()
