#! /usr/bin/python -u
#
# Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# This software product is a proprietary product of Nvidia Corporation
# and its affiliates (the "Company") and all right, title, and interest
# in and to the software product, including all associated intellectual
# property rights, are and shall remain exclusively with the Company.
#
# This software product is governed by the End User License Agreement
# provided with the software product.
# All Rights reserved.
#
# This program is a Net-SNMP pass persist script that services the
# Nvidia Cumulus(R) Linux(R) BGP unumbered functionality RFC 5549
# Derived from bgp4_pp.py
#
# To activate this script, snmpd must be running and this
# script should be installed /usr/share/snmp/bgpun_pp.py
# and include the following line in /etc/snmp/snmpd.conf
#
#   pass_persist .1.3.6.1.4.1.40310.7 /usr/share/snmp/vrf_bgpun_pp.py
#
# Corresponding MIB is located at 
#   /usr/share/snmp/mibs/Cumulus-BGPVRF-MIB.txt

import subprocess
import syslog
import json
import traceback
import snmp_passpersist as snmp
import sys
import getopt
import re
import netaddr
import socket
import struct

# The Nvidia Cumulus Linux custom VRF BGP entries starts with this object ID
oid_base = '.1.3.6.1.4.1.40310.7'
# by default do show path table unless user asks for it
pp = snmp.PassPersist(oid_base)
showPathTable = False
refreshInterval = 60

vrf_re = re.compile('^vrf\s(\S+)\sid\s(\d+)\stable\s(\S+).*\n', re.M)

dfltVrfTable = 254
dfltVrfTableIdentity = 0

idTypeInterface = 0
idTypeIPv4 = 1
idTypeIPv6 = 2

short_options = "di:p"
long_options = ["debug", "interval=", "include-paths"]

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',
                    'bgpPeerMinRouteAdvertisementInterval',
                    'bgpPeerInUpdateElapsedTime',
                    'bgpPeerIface','bgpPeerDesc','bgpPeerIfindex','bgpPeerAddressType']

bgpPeerEntryList = list(enumerate(bgpPeerEntryRows, start=1))

# construct the peer entry table dictionary
bgpPeerEntry = {'bgpPeerIdentifier' : {'oid' : 1, 'type' : pp.add_ip,
                                       'jsonName' : ['remoteRouterId'], 'default' : '0.0.0.0'},
            'bgpPeerState' : {'oid' : 2, 'type' : pp.add_int,
                              'jsonName' : ['bgpState'], 'default' : 1},
            'bgpPeerAdminStatus' : {'oid' : 3, 'type' : pp.add_int,
                                    'jsonName' : ['adminShutDown'], 'default' : 2},
            'bgpPeerNegotiatedVersion' : {'oid' : 4, 'type' : pp.add_int,
                                          'jsonName' : ['bgpVersion'],
                                          'default' : 0},
            'bgpPeerLocalAddr' : {'oid' : 5, 'type' : pp.add_str,
                                  'jsonName' : ['hostLocal'],
                                  'default' : '0.0.0.0'},
            'bgpPeerLocalPort' : {'oid' : 6, 'type' : pp.add_int,
                                  'jsonName' : ['portLocal'], 'default' : 0},
            'bgpPeerRemoteAddr' : {'oid' : 7, 'type' : pp.add_str,
                                   'jsonName' : ['hostForeign'],
                                   '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' : ['connectionsEstablished'],
                                                  'default' : 0},
            'bgpPeerFsmEstablishedTime' : {'oid' : 16, 'type' : pp.add_gau,
                                           'jsonName' : ['bgpTimerUpMsec'],
                                           '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},
            'bgpPeerMinRouteAdvertisementInterval' : {'oid' : 23, 'type' : pp.add_int,
                          'jsonName' : ['minBtwnAdvertisementRunsTimerMsecs'],
                          'default' : 0},
            'bgpPeerInUpdateElapsedTime' : {'oid' : 24, 'type' : pp.add_gau,
                          'jsonName' : ["bgpInUpdateElapsedTimeMsecs"],
                          'default' : 0},
            'bgpPeerIface' : {'oid' : 25, 'type' : pp.add_str,
                          'jsonName' : [],
                          'default' : "iface"},
            'bgpPeerDesc' : {'oid' : 26, 'type' : pp.add_str,
                          'jsonName' : ["nbrDesc"],
                          'default' : "NA"},
            'bgpPeerIfindex' : {'oid' : 27, 'type' : pp.add_int,
                             'jsonName' : [],
                             'default' : 0},
            'bgpPeerAddressType' : {'oid' : 28, 'type' : pp.add_int,
                             'jsonName' : [],
                             '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_str,
                                       'jsonName' : ['peerId'],
                                       'default' : '::'},
                 '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' : ['path'],
                                                'default' : '00 00'},
                 'bgp4PathAttrNextHop' : {'oid' : 6, 'type' : pp.add_str,
                                          'jsonName' : ['nexthops','ip'],
                                          'default' : '::'},
                 'bgp4PathAttrMultiExitDisc' : {'oid' : 7, 'type' : pp.add_int,
                                                'jsonName' : ['med'],
                                                'default' : -1},
                 'bgp4PathAttrLocalPref' : {'oid' : 8, 'type' : pp.add_int,
                                            'jsonName' : ['locPrf'],
                                            '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' : ['locPrf'],
                                                '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 in ['bgpPeerNegotiatedVersion', 'bgpPeerIdentifier']:
        # only show the peer and 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 == 'bgpPeerFsmEstablishedTime':
        # The time established key 'bgpTimerUpMsec' only exists while in the established state.
        # If we are not currently established but have been previously, we need to use the key 'lastResetTimerMsecs'.
        # If we have never been established, based on the MIB definition, we need to return 0.
        value = int(traverse(peerDict, jsonList, default) or 0)
        if not (value > 0 or peerDict.get('connectionsEstablished', 0) == 0):
            value = peerDict.get('lastResetTimerMsecs', 0)
        value /= 1000
    elif rowname in [
        '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 rowname == 'bgpPeerAdminStatus':
        # The key is only present if the peer is actually shutdown and it's a bool.
        # So if getValue above returns True, reset the value to 1. Otherwise the default of 2 means the key was
        # not present.
        value = 1 if traverse(peerDict, jsonList, default) is True else default
    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 per_vrf_update(name, table, identity):
    vrf = name
    try:
        subprocess.check_output(['/usr/bin/pgrep', 'bgpd'])
        ipBgpNeig = get_json_output(commandList=['sudo','vtysh','-c','show ip bgp vrf {} neighbors json'.format(vrf)])
        ipBgpSummary  = get_json_output(commandList=['sudo','vtysh','-c','show ip bgp vrf {} summary json'.format(vrf)])
    except Exception as e:
        return

    peerList = ipBgpNeig.keys()
    ipv4PeerList = []
    ifacePeerList = []
    iface_dict = {}
    id_type = {}
    try:
        pp.add_int('6.{}'.format(table), table)
        pp.add_int('7.{}'.format(table), identity)
        pp.add_str('8.{}'.format(table), name)
    except Exception as e:
        syslog.syslog('Error: %s' % e)

    if not (ipBgpNeig or ipBgpSummary): 
        return

    for peer in peerList:
        try:
            ifacePeerList.append(peer)
        except ValueError:
            continue

    bgpSummary = ipBgpSummary.get('ipv4Unicast')

    for if_peer in ifacePeerList:
        try:
            if bgpSummary['peers'].get(if_peer, {}).get('idType') == 'interface':
                ipOutput = get_json_output(commandList=['ip', '-j', 'link', 'show', '{}'.format(if_peer)])
                ifindexIp = str(netaddr.IPAddress(ipOutput[0].get('ifindex')))

                ipv4PeerList.append(ifindexIp)
                iface_dict[ifindexIp] = if_peer
                id_type[ifindexIp] = idTypeInterface
            else:
                neigh_host = ipBgpNeig[if_peer].get('hostForeign', if_peer)
                ipv4PeerList.append(neigh_host)
                iface_dict[neigh_host] = if_peer
                if bgpSummary['peers'].get(if_peer, {}).get('idType') == 'ipv4':
                    id_type[neigh_host] = idTypeIPv4
                else :
                    id_type[neigh_host] = idTypeIPv6
        except Exception as e:
            syslog.syslog('Error: EXCEPTION=%s' % (e))
            continue

    try:
        bgpLocalAs = ipBgpNeig[ifacePeerList[0]]['localAs']
    except:
        bgpLocalAs = 0

    ##################### bgpPeerEntryTable ####################################
    bgpPeerEntryTable = "3.1.1"
    for rowval, rowname in bgpPeerEntryList:
        # show all peers for each row
        for peer in ipv4PeerList:
            jsonList =  bgpPeerEntry[rowname]['jsonName']
            default =  bgpPeerEntry[rowname]['default']
            oid = bgpPeerEntry[rowname].get('oid', rowval)
            peerDict = ipBgpNeig[iface_dict[peer]]
            state = peerDict['bgpState']

            if id_type[peer] == idTypeIPv6:
                octets = re.split(':', peer.replace('::',':'))
                dotted_peer = str(int(octets[0],16)) + "".join('.'+str(int(octet, 16)) for octet in octets[1:])
                newOid = "%s.%s.%s.%s.%s" % (bgpPeerEntryTable, oid, table, id_type[peer], dotted_peer)
            else:
                newOid = "%s.%s.%s.%s.%s" % (bgpPeerEntryTable, oid, table, id_type[peer], peer)

            myval = getValue(peer=peer, rowname=rowname, state=state, default=default,
                             jsonList=jsonList, peerDict=peerDict)
            if id_type[peer] == idTypeInterface:
                if rowname == 'bgpPeerIface':
                    myval = iface_dict[peer]
                if rowname == 'bgpPeerIfindex':
                    myval = struct.unpack("!I", socket.inet_aton(peer))[0]
            if rowname == 'bgpPeerAddressType':
                myval = id_type[peer]

            func = bgpPeerEntry[rowname]['type']
            func(newOid, myval)

    bgpVersion = '10'
    try:
        pp.add_oct('1.{}'.format(table), bgpVersion)
        pp.add_int('2.{}'.format(table), bgpLocalAs)
        bgpIdentifier = bgpSummary.get('routerId') 
        pp.add_ip('4.{}'.format(table), bgpIdentifier)

    except Exception, err:
        syslog.syslog('leaving....Exception=%s' % err)
        syslog.syslog('traceback=%s' % traceback.format_exc())

    ##################### bgpPathEntryTable ####################################
    # we only update and show the path table if the user configured it
    if showPathTable == False:
        return
    try:
        subprocess.check_output(['/usr/bin/pgrep', 'bgpd'])
        ipBgpPath = get_json_output(commandList=['sudo','vtysh','-c','show ip bgp vrf {} json'.format(vrf)])
    except Exception as e:
        return

    if not ipBgpPath:
        return

    # the bgpRcvdPathAttrTable (5) is obsolete so we do not show it
    bgpPathAttrTable = '5'
    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]))
    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)]
            count = 1
            for path in routeList:
                path['prefixlen'] = prefix
                path['prefix'] = route
                newOid = "%s.%s.%s.%s.%s.%s.%s" % \
                         (bgpPathAttrTable, bgpPathEntryTable, rowval, table, route, prefix,
                          count)
                count = count + 1
                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) / 0x100))
                            buf = buf + (' %02x' % (int(path) & 0xff))
                        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)
##############################################################################
def update():
    #To handle bgp unnumbered configuration on default vrf
    per_vrf_update("default", dfltVrfTable, dfltVrfTableIdentity)

    try:
        cmd = ['sudo', 'vtysh', '-c', 'show vrf']
        vrfsList = vrf_re.findall(subprocess.check_output(cmd,shell=False))
        for vrf in vrfsList:
            table = vrf[2]
            name = vrf[0]
            identity = vrf[1]
            per_vrf_update(name, table, identity)
    except Exception as e:
        syslog.syslog('Error: command %s EXCEPTION=%s' % (' '.join(cmd), e))
        return

pp.debug = False
syslog.syslog("Starting BGP VRF SNMP script...")

full_cmd_arguments = sys.argv
argument_list = full_cmd_arguments[1:]

try:
    arguments, values = getopt.getopt(argument_list, short_options, long_options)
except getopt.error as err:
    syslog.syslog(str(err))
    sys.exit(2)

for current_argument, current_value in arguments:
    if current_argument in ("-d", "--debug"):
        pp.debug = True
    elif current_argument in ("-p", "--include-paths"):
        showPathTable = True
    elif current_argument in ("-i", "--interval"):
        refreshInterval = float(current_value)

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