#! /usr/bin/python -u
# Copyright 2016 Cumulus Networks LLC, all rights reserved
#
# Implementation of the BGP4 MIB (RFC 4273) which includes
# the bgpPeerTable, bgp4PathAttrTable.
#
# To activate this script, snmpd must be running and this
# script should be installed /usr/share/snmp/bgp4_pp.py
# and include the following line in /etc/snmp/snmpd.conf
#
#   pass_persist .1.3.6.1.2.1.15  /usr/share/snmp/bgp4_pp.py
#

import subprocess
import sys
import json
import syslog
import ipaddr, struct, socket
import snmp_passpersist as snmp

BGP4_MIB = '.1.3.6.1.2.1.15'
pp = snmp.PassPersist(BGP4_MIB)

# by default do not show path table unless user asks for it
showPathTable = False

BGP4_MIB = '.1.3.6.1.2.1.15'

peerstate = {'Idle':1,
             'Connect':2,
             'Active':3,
             'OpenSent':4,
             'OpenConfirm':5,
             'Established':6}

attrorigin = {'IGP':1,
              'EGP':2,
              'incomplete':3}
bgpPeerEntryRows = ['bgpPeerIdentifier', 'bgpPeerState', 'bgpPeerAdminStatus',
                    'bgpPeerNegotiatedVersion', 'bgpPeerLocalAddr', 'bgpPeerLocalPort',
                    'bgpPeerRemoteAddr', 'bgpPeerRemotePort', 'bgpPeerRemoteAs',
                    'bgpPeerInUpdates', 'bgpPeerOutUpdates', 'bgpPeerInTotalMessages',
                    'bgpPeerOutTotalMessages', 'bgpPeerLastError',
                    'bgpPeerFsmEstablishedTransitions', 'bgpPeerFsmEstablishedTime',
                    'bgpPeerConnectRetryInterval', 'bgpPeerHoldTime', 'bgpPeerKeepAlive',
                    'bgpPeerHoldTimeConfigured', 'bgpPeerKeepAliveConfigured',
                    'bgpPeerMinASOriginationInterval',
                    'bgpPeerMinRouteAdvertisementInterval',
                    'bgpPeerInUpdateElapsedTime']
bgpPeerEntryList = list(enumerate(bgpPeerEntryRows, start=1))

# construct the peer entry table dictionary
bgpPeerEntry = {'bgpPeerIdentifier' : {'oid' : 1, 'type' : pp.add_ip,
                                       'jsonName' : [], 'default' : '0.0.0.0'},
            'bgpPeerState' : {'oid' : 2, 'type' : pp.add_int,
                              'jsonName' : ['bgpState'], 'default' : 1},
            'bgpPeerAdminStatus' : {'oid' : 3, 'type' : pp.add_int,
                                    'jsonName' : [], 'default' : 2},
            'bgpPeerNegotiatedVersion' : {'oid' : 4, 'type' : pp.add_int,
                                          'jsonName' : ['bgpVersion'],
                                          'default' : 0},
            'bgpPeerLocalAddr' : {'oid' : 5, 'type' : pp.add_ip,
                                  'jsonName' : ['hostLocal'],
                                  'default' : '0.0.0.0'},
            'bgpPeerLocalPort' : {'oid' : 6, 'type' : pp.add_int,
                                  'jsonName' : ['portLocal'], 'default' : 0},
            'bgpPeerRemoteAddr' : {'oid' : 7, 'type' : pp.add_ip,
                                   'jsonName' : ['remoteRouterId'],
                                   'default' : '0.0.0.0'},
            'bgpPeerRemotePort' : {'oid' : 8, 'type' : pp.add_int,
                                   'jsonName' : ['portForeign'],
                                   'default' : 0},
            'bgpPeerRemoteAs' : {'oid' : 9, 'type' : pp.add_int,
                                 'jsonName' : ['remoteAs'], 'default' : 0},
            'bgpPeerInUpdates' : {'oid' : 10, 'type' : pp.add_cnt_32bit,
                                  'jsonName' : ["messageStats", "updatesRecv"],
                                  'default' : 0},
            'bgpPeerOutUpdates' : {'oid' : 11, 'type' : pp.add_cnt_32bit,
                                   'jsonName' : ["messageStats", "updatesSent"],
                                   'default' : 0},
            'bgpPeerInTotalMessages' : {'oid' : 12, 'type' : pp.add_cnt_32bit,
                                        'jsonName' : ["messageStats", "totalRecv"],
                                        'default' : 0},
            'bgpPeerOutTotalMessages' : {'oid' : 13, 'type' : pp.add_cnt_32bit,
                                         'jsonName' : ["messageStats", "totalSent"],
                                         'default' : 0},
            'bgpPeerLastError' : {'oid' : 14, 'type' : pp.add_oct,
                                  'jsonName' : ['lastErrorCodeSubcode'],
                                  'default' : '00 00'},
            'bgpPeerFsmEstablishedTransitions' : {'oid' : 15,
                                                  'type' : pp.add_cnt_32bit,
                                  'jsonName' : ['addressFamilyInfo',
                                                'connectionsEstablished'],
                                                  'default' : 0},
            'bgpPeerFsmEstablishedTime' : {'oid' : 16, 'type' : pp.add_gau,
                                           'jsonName' : ['bgpTimerUp'],
                                           'default' : 0},
            'bgpPeerConnectRetryInterval' : {'oid' : 17, 'type' : pp.add_int,
                                             'jsonName' : ['connectRetryTimer'],
                                             'default' : 0},
            'bgpPeerHoldTime' : {'oid' : 18, 'type' : pp.add_int,
                                 'jsonName' : ['bgpTimerHoldTimeMsecs'],
                                 'default' : 0},
            'bgpPeerKeepAlive' : {'oid' : 19, 'type' : pp.add_int,
                                  'jsonName' : ['bgpTimerKeepAliveIntervalMsecs'],
                                  'default' : 0},
            'bgpPeerHoldTimeConfigured' : {'oid' : 20, 'type' : pp.add_int,
                          'jsonName' : ['bgpTimerConfiguredHoldTimeMsecs'],
                          'default' : 0},
            'bgpPeerKeepAliveConfigured' : {'oid' : 21, 'type' : pp.add_int,
                          'jsonName' : ['bgpTimerConfiguredKeepAliveIntervalMsecs'],
                          'default' : 0},
            'bgpPeerMinASOriginationInterval' : {'oid' : 22, 'type' : pp.add_int,
                                                 'jsonName' : [], 'default' : 0},
            'bgpPeerMinRouteAdvertisementInterval' : {'oid' : 23, 'type' : pp.add_int,
                          'jsonName' : ['minBtwnAdvertisementRunsTimerMsecs'],
                          'default' : 0},
            'bgpPeerInUpdateElapsedTime' : {'oid' : 24, 'type' : pp.add_gau,
                          'jsonName' : ["bgpInUpdateElapsedTimeMsecs"],
                          'default' : 0}}

bgp4PathEntryRows = ['bgp4PathAttrPeer', 'bgp4PathAttrIpAddrPrefixLen',
                     'bgp4PathAttrIpAddrPrefix',
                     'bgp4PathAttrOrigin', 'bgp4PathAttrASPathSegment',
                     'bgp4PathAttrNextHop',
                     'bgp4PathAttrMultiExitDisc', 'bgp4PathAttrLocalPref',
                     'bgp4PathAttrAtomicAggregate',
                     'bgp4PathAttrAggregatorAS', 'bgp4PathAttrAggregatorAddr',
                     'bgp4PathAttrCalcLocalPref',
                     'bgp4PathAttrBest', 'bgp4PathAttrUnknown']

bgp4PathEntryList = list(enumerate(bgp4PathEntryRows, start=1))
bgp4PathEntry = {'bgp4PathAttrPeer' : {'oid' : 1, 'type' : pp.add_ip,
                                       'jsonName' : ['peerId'],
                                       'default' : '0.0.0.0'},
                 'bgp4PathAttrIpAddrPrefixLen' : {'oid' : 2, 'type' : pp.add_int,
                                                  'jsonName' : ['prefixlen'],
                                                  'default' : 0},
                 'bgp4PathAttrIpAddrPrefix' : {'oid' : 3, 'type' : pp.add_ip,
                                               'jsonName' : ['prefix'],
                                               'default' : '0.0.0.0'},
                 'bgp4PathAttrOrigin' : {'oid' : 4, 'type' : pp.add_int,
                                         'jsonName' : ['origin'], 'default' : 3},
                 'bgp4PathAttrASPathSegment' : {'oid' : 5, 'type' : pp.add_oct,
                                                'jsonName' : ['aspath'],
                                                'default' : '00 00'},
                 'bgp4PathAttrNextHop' : {'oid' : 6, 'type' : pp.add_ip,
                                          'jsonName' : ['nexthops','ip'],
                                          'default' : '0.0.0.0'},
                 'bgp4PathAttrMultiExitDisc' : {'oid' : 7, 'type' : pp.add_int,
                                                'jsonName' : ['med'],
                                                'default' : -1},
                 'bgp4PathAttrLocalPref' : {'oid' : 8, 'type' : pp.add_int,
                                            'jsonName' : ['localpref'],
                                            'default' : -1},
                 'bgp4PathAttrAtomicAggregate' : {'oid' : 9, 'type' : pp.add_int,
                                                  'jsonName' : [],
                                                  'default' : 1},
                 'bgp4PathAttrAggregatorAS' : {'oid' : 10, 'type' : pp.add_int,
                                               'jsonName' : ["aggregatorAs"],
                                               'default' : 0},
                 'bgp4PathAttrAggregatorAddr' : {'oid' : 11, 'type' : pp.add_ip,
                                                 'jsonName' : ["aggregatorId"],
                                                 'default' : '0.0.0.0'},
                 'bgp4PathAttrCalcLocalPref' : {'oid' : 12, 'type' : pp.add_int,
                                                'jsonName' : ['localpref'],
                                                'default' : -1},
                 'bgp4PathAttrBest' : {'oid' : 13, 'type' : pp.add_int,
                                       'jsonName' : ["bestpath"], 'default' : 1},
                 'bgp4PathAttrUnknown' : {'oid' : 14, 'type' : pp.add_str,
                                          'jsonName' : [], 'default' : ''}
                 }

def traverse(obj, key, default):
    '''
    recursive func where obj is a dictionary and key is a list.
    we need to grab subdictionaries and remove a key
    if we get an obj as a list, then just grab the first one since we
    are only dealing with IPv4 (IPv6 can have multiple nexthops (linklocal).
    '''
    if len(key) <= 1:
        # we only return if we have one key left
        try:
            if type(obj) is list:
                obj = obj[0]
            value = obj[key[0]]
        except (KeyError, IndexError):
            value = default
        return(value)
    return(traverse(obj[key[0]], key[1:], default))

def getValue(peer=None, rowname=None, state=None, default=None,
             jsonList=None, peerDict=None):
    '''
    Handle getting the actual value as this can vary depending on
    state and the actual row we are trying to get.
    '''
    if rowname == 'bgpPeerIdentifier':
        # we cannot show the actual peer identifier unless the state is one of these
        if state == 'Established' or state == 'OpenConfirm':
            value = peer
        else:
            value = default
    elif rowname == 'bgpPeerNegotiatedVersion':
        # only show the negotiated version if we are established or openconfirm
        if state == 'Established' or state == 'OpenConfirm':
            value = traverse(peerDict, jsonList, default)
        else:
            value = default
    elif rowname == 'bgpPeerLastError':
        # the last error must be 4 character hex string with first two being error code
        # and second two being error subcode.  This may or may not exist.
        val = traverse(peerDict, jsonList, default) or '00 00'
        if len(val) != 4:
            val = '00 00'
        value = '%s %s' % (val[:2], val[2:])
    elif rowname == 'bgpPeerState':
        # state is an integer not a string
        value = peerstate.get(traverse(peerDict, jsonList, default), 0)
    elif rowname in ['bgpPeerFsmEstablishedTime', 'bgpPeerHoldTime', 'bgpPeerKeepAlive',
                     'bgpPeerHoldTimeConfigured', 'bgpPeerKeepAliveConfigured',
                     'bgpPeerMinRouteAdvertisementInterval', 'bgpPeerInUpdateElapsedTime']:
        # time was given to us in ms, convert to seconds
        value = int(traverse(peerDict, jsonList, default) or 0)/1000
    elif jsonList == []:
        value = default
    else:
        # for all else, just get it
        value = traverse(peerDict, jsonList, default) or 0
    return(value)

def get_json_output(commandList=None):
    '''
    This grabs the JSON output of a CLI command and returns an array.
    '''
    global syslog_already_logged
    outArray = {}
    if (not commandList) or (commandList == []):
        syslog.syslog('Error: called get_json_output with'
                      ' invalid command=%s' % \
                      commandList)
        syslog_already_logged = True
        return {}
    try:
        outArray = json.loads(subprocess.check_output((commandList),
                                                      shell=False),
                              encoding="latin-1")
    except Exception as e:
        outArray = {}
        syslog.syslog('Error: command %s EXCEPTION=%s' % \
                      (' '.join(commandList), e))
        syslog_already_logged = True
    return outArray

def update():
    """
    Simply grab the output of vtysh commands and stick them in our hashed array
    """
    # version is a vector of supported protocol versions with MSB of first octet
    # as bit 0 and version is i+1 if bit i is set.  We hardcode this to version 4.
    bgpVersion = '10'
    # return a hex string
    pp.add_oct("1", bgpVersion)
    # grab an array of neighbor entries
    # we have created showpeers and showpaths to simplify sudoers.d/snmp
    ipBgpNeig = get_json_output(commandList=['sudo', '/usr/share/snmp/showpeers'])
    ipBgpSummary  = get_json_output(commandList=['sudo', '/usr/share/snmp/showsummary'])

    peerList = ipBgpNeig.keys()
    ipv4PeerList = []
    for peer in peerList:
        try:
            if ipaddr.IPNetwork(peer).version == 4:
                ipv4PeerList.append(peer)
        except ValueError:
            # ignore IPv6 since this rfc cannot handle it
            continue
    try:
        bgpLocalAs = ipBgpNeig[ipv4PeerList[0]]['localAs']
    except:
        bgpLocalAs = 0
    pp.add_int("2.0", bgpLocalAs)
    ##################### bgpPeerEntryTable ####################################
    bgpPeerEntryTable = "3.1"
    # peers should be sorted by ip address because snmp expects it.
    ipv4PeerList = sorted(ipv4PeerList,
                          key=lambda ip: struct.unpack("!L", socket.inet_aton(ip))[0])
    for rowval, rowname in bgpPeerEntryList:
        # show all peers for each row
        for peer in ipv4PeerList:
            newOid = "%s.%s.%s" % (bgpPeerEntryTable, rowval, peer)
            jsonList =  bgpPeerEntry[rowname]['jsonName']
            default =  bgpPeerEntry[rowname]['default']
            peerDict = ipBgpNeig[peer]
            state = peerDict['bgpState']
            myval = getValue(peer=peer, rowname=rowname, state=state, default=default,
                             jsonList=jsonList, peerDict=peerDict)
            func = bgpPeerEntry[rowname]['type']
            func(newOid, myval)

    # local system identifier IP address from  vtysh -c "show ip bgp summary json"
    bgpIdentifier = ipBgpSummary.get('routerId','0.0.0.0')
    pp.add_ip('4', bgpIdentifier )

    ##################### bgpPathEntryTable ####################################
    # we only update and show the path table if the user configured it
    if not showPathTable:
        return

    ipBgpPath  = get_json_output(commandList=['sudo', '/usr/share/snmp/showpaths'])

    # the bgpRcvdPathAttrTable (5) is obsolete so we do not show it
    bgpPathAttrTable = '6'
    bgpPathEntryTable = '1'
    pathList = ipBgpPath.get('routes', {}).keys()
    if pathList == []:
        return
    pathList = [(x.split('/')[0], (x.split('/')[1])) for x in pathList]
    # now just sort the routes based on the route
    mykey = lambda ip: struct.unpack("!L", socket.inet_aton(ip[0]))
    pathList = sorted(pathList, key=lambda ip: struct.unpack("!L", socket.inet_aton(ip[0]))[0])
    for rowval, rowname in bgp4PathEntryList:
        for route,prefix in pathList:
            # show all paths for each row
            path = '%s/%s' % (route,prefix)
            routeList = ipBgpPath['routes']['%s/%s' % (route,prefix)]
            routeList = sorted(routeList,
                   key=lambda ip: struct.unpack("!L", socket.inet_aton(ip['peerId'])))
            for path in routeList:
                path['prefixlen'] = prefix
                path['prefix'] = route
                newOid = "%s.%s.%s.%s.%s.%s" % \
                         (bgpPathAttrTable, bgpPathEntryTable, rowval, route, prefix,
                          path['peerId'])
                default = bgp4PathEntry[rowname]['default']
                value = traverse(path, bgp4PathEntry[rowname]['jsonName'], default)
                if rowname == 'bgp4PathAttrOrigin':
                    value = attrorigin[value]
                elif rowname == 'bgp4PathAttrBest':
                    if path.get('bestpath', False):
                        # true is 2
                        value = 2
                    else:
                        value = 1
                elif rowname == 'bgp4PathAttrASPathSegment':
                    # the attribute AS path requires a TLV binary string
                    # handle only AS_SEQUENCE types (2)
                    # the paths are space seperated AS numbers
                    if value:
                        aspaths = value.strip().split()
                        buf = '02 %02x' % (len(aspaths))
                        for path in aspaths:
                            buf = buf + (' %02x' % int(path))
                        value = buf

                if not value:
                    value = bgp4PathEntry[rowname]['default']
                # now return the correct type and value
                func = bgp4PathEntry[rowname]['type']
                # now print the best path info for this route
                func(newOid, value)
    return

################################################################################
if len(sys.argv) > 1 and sys.argv[1] == '--include-paths':
    showPathTable = True

syslog.syslog("starting...")

pp.debug = False

if pp.debug:
    update()
else:
    pp.start(update, 60)
