From f3a6c3519a43919391d1aebd420f555cb2c6a93f Mon Sep 17 00:00:00 2001 From: IWASE Yusuke Date: Fri, 22 Dec 2017 09:32:44 +0900 Subject: test/lib: Add YABGP container Note: This implementation supposes the YABGP version 0.4.0. Signed-off-by: IWASE Yusuke --- test/lib/yabgp_helper.py | 191 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 191 insertions(+) create mode 100644 test/lib/yabgp_helper.py (limited to 'test/lib/yabgp_helper.py') diff --git a/test/lib/yabgp_helper.py b/test/lib/yabgp_helper.py new file mode 100644 index 00000000..287daeaa --- /dev/null +++ b/test/lib/yabgp_helper.py @@ -0,0 +1,191 @@ +#!/usr/bin/env python + +from __future__ import absolute_import +from __future__ import print_function + +import json +import logging +import sys + +from flask import Blueprint +import flask + +from yabgp.agent import prepare_service +from yabgp.common import constants +from yabgp.handler import BaseHandler +from yabgp.api import app +from yabgp.api import utils as api_utils +from yabgp.api.v1 import auth + +LOG = logging.getLogger(__name__) +blueprint = Blueprint('v1-ext', __name__) +_send_update = api_utils.send_update + + +def _extract_afi_safi(attr): + # MP_REACH_NLRI(14) or MP_UNREACH_NLRI(15) + nlri = attr.get(14, None) or attr.get(15, None) + if nlri and 'afi_safi' in nlri: + afi_safi = nlri.get('afi_safi', None) + return constants.AFI_SAFI_DICT.get(tuple(afi_safi), None) + return 'ipv4' + + +def _construct_msg(attr, nlri, withdraw): + return { + 'afi_safi': _extract_afi_safi(attr), + 'attr': attr, + 'nlir': nlri or [], + 'withdraw': withdraw or [], + } + + +def _extract_nlri_list(msg): + try: + if msg['afi_safi'] == 'ipv4': + return msg['nlri'] + return msg['attr'][14]['nlri'] # MP_REACH_NLRI(14) + except KeyError: + # Ignore case when no nlri + pass + return [] + + +def _extract_withdraw_list(msg): + try: + if msg['afi_safi'] == 'ipv4': + return msg['withdraw'] + return msg['attr'][15]['withdraw'] # MP_UNREACH_NLRI(15) + except KeyError: + # Ignore case when no withdraw + pass + return [] + + +def _nlri_key(nlri): + if isinstance(nlri, (dict, list)): + return json.dumps(nlri) + return str(nlri) + + +# ADJ_RIB_IN = { +# : { +# : { +# : , +# ... +# }, +# ... +# }, +# ... +# } +ADJ_RIB_IN = {} + + +@blueprint.route('/peer//adj-rib-in') +@auth.login_required +@api_utils.log_request +def peer_adj_rib_in(peer_ip): + """ + Dumps one peer's adj-RIB-in. + """ + return flask.jsonify(ADJ_RIB_IN.get(peer_ip, {})) + + +# ADJ_RIB_OUT = { +# : { +# : { +# : , +# ... +# }, +# ... +# }, +# ... +# } +ADJ_RIB_OUT = {} + + +@blueprint.route('/peer//adj-rib-out') +@auth.login_required +@api_utils.log_request +def peer_adj_rib_out(peer_ip): + """ + Dumps one peer's adj-RIB-out. + """ + return flask.jsonify(ADJ_RIB_OUT.get(peer_ip, {})) + + +app.app.register_blueprint(blueprint, url_prefix='/v1-ext') + + +def send_update(peer_ip, attr, nlri, withdraw): + """ + Wrapper of "yabgp.api.send_update" in order to hook sending UPDATE + messages via REST API. + """ + msg = _construct_msg(attr, nlri, withdraw) + afi_safi = msg['afi_safi'] + ADJ_RIB_OUT.setdefault(peer_ip, {}) + ADJ_RIB_OUT[peer_ip].setdefault(afi_safi, {}) + rib = ADJ_RIB_OUT[peer_ip][afi_safi] + for _nlri in _extract_nlri_list(msg): + rib[_nlri_key(_nlri)] = msg + for _withdraw in _extract_withdraw_list(msg): + rib.pop(_nlri_key(_withdraw), None) + return _send_update(peer_ip, attr, nlri, withdraw) + + +setattr(api_utils, 'send_update', send_update) + + +class CliHandler(BaseHandler): + + def __init__(self): + super(CliHandler, self).__init__() + + def init(self): + pass + + def on_update_error(self, peer, timestamp, msg): + LOG.info('[-] UPDATE ERROR: %s', msg) + + def route_refresh_received(self, peer, msg, msg_type): + LOG.info('[+] ROUTE_REFRESH received: %s', msg) + + def keepalive_received(self, peer, timestamp): + LOG.debug('[+] KEEPALIVE received: %s', peer.factory.peer_addr) + + def open_received(self, peer, timestamp, result): + LOG.info('[+] OPEN received: %s', result) + + def update_received(self, peer, timestamp, msg): + LOG.info('[+] UPDATE received: %s', msg) + peer_addr = peer.factory.peer_addr + afi_safi = msg['afi_safi'] + ADJ_RIB_IN.setdefault(peer.factory.peer_addr, {}) + ADJ_RIB_IN[peer_addr].setdefault(afi_safi, {}) + rib = ADJ_RIB_IN[peer_addr][afi_safi] + for nlri in _extract_nlri_list(msg): + rib[_nlri_key(nlri)] = msg + for withdraw in _extract_withdraw_list(msg): + rib.pop(_nlri_key(withdraw), None) + + def notification_received(self, peer, msg): + LOG.info('[-] NOTIFICATION received: %s', msg) + + def on_connection_lost(self, peer): + LOG.info('[-] CONNECTION lost: %s', peer.factory.peer_addr) + + def on_connection_failed(self, peer, msg): + LOG.info('[-] CONNECTION failed: %s', msg) + + +def main(): + try: + cli_handler = CliHandler() + prepare_service(handler=cli_handler) + except Exception as e: + print(e) + + +if __name__ == '__main__': + sys.exit(main()) -- cgit v1.2.3