test/lib: Add YABGP container

Note: This implementation supposes the YABGP version 0.4.0.

Signed-off-by: IWASE Yusuke <iwase.yusuke0@gmail.com>

3 files changed, 694 insertions, 0 deletions

diff --git a/test/lib/base.py b/test/lib/base.py
index 19452896..c7afba6d 100644
--- a/test/lib/base.py
+++ b/test/lib/base.py
@@ -48,6 +48,33 @@ BGP_ATTR_TYPE_CLUSTER_LIST = 10
 BGP_ATTR_TYPE_MP_REACH_NLRI = 14
 BGP_ATTR_TYPE_EXTENDED_COMMUNITIES = 16
 
+FLOWSPEC_NAME_TO_TYPE = {
+    "destination": 1,
+    "source": 2,
+    "protocol": 3,
+    "port": 4,
+    "destination-port": 5,
+    "source-port": 6,
+    "icmp-type": 7,
+    "icmp-code": 8,
+    "tcp-flags": 9,
+    "packet-length": 10,
+    "dscp": 11,
+    "fragment": 12,
+    "label": 13,
+    "ether-type": 14,
+    "source-mac": 15,
+    "destination-mac": 16,
+    "llc-dsap": 17,
+    "llc-ssap": 18,
+    "llc-control": 19,
+    "snap": 20,
+    "vid": 21,
+    "cos": 22,
+    "inner-vid": 23,
+    "inner-cos": 24,
+}
+
 # with this label, we can do filtering in `docker ps` and `docker network prune`
 TEST_CONTAINER_LABEL = 'gobgp-test'
 TEST_NETWORK_LABEL = TEST_CONTAINER_LABEL
diff --git a/test/lib/yabgp.py b/test/lib/yabgp.py
new file mode 100644
index 00000000..c3e5e0ac
--- /dev/null
+++ b/test/lib/yabgp.py
@@ -0,0 +1,476 @@
+# Copyright (C) 2017 Nippon Telegraph and Telephone Corporation.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import print_function
+
+import json
+import os
+import time
+
+from fabric import colors
+from fabric.api import local
+from fabric.utils import indent
+
+from lib.base import (
+    FLOWSPEC_NAME_TO_TYPE,
+    BGPContainer,
+    CmdBuffer,
+    try_several_times,
+)
+
+
+class YABGPContainer(BGPContainer):
+
+    WAIT_FOR_BOOT = 1
+    SHARED_VOLUME = '/etc/yabgp'
+
+    def __init__(self, name, asn, router_id,
+                 ctn_image_name='osrg/yabgp:v0.4.0'):
+        super(YABGPContainer, self).__init__(name, asn, router_id,
+                                             ctn_image_name)
+        self.shared_volumes.append((self.config_dir, self.SHARED_VOLUME))
+
+    def _copy_helper_app(self):
+        import lib
+        mod_dir = os.path.dirname(lib.__file__)
+        local('docker cp {0}/yabgp_helper.py'
+              ' {1}:/root/'.format(mod_dir, self.name))
+
+    def _start_yabgp(self):
+        self.local(
+            'python /root/yabgp_helper.py'
+            ' --config-file {0}/yabgp.ini'.format(self.SHARED_VOLUME),
+            detach=True)
+
+    def run(self):
+        super(YABGPContainer, self).run()
+        # self.create_config() is called in super class
+        self._copy_helper_app()
+        self._start_yabgp()
+        return self.WAIT_FOR_BOOT
+
+    def create_config(self):
+        # Currently, supports only single peer
+        c = CmdBuffer('\n')
+        c << '[DEFAULT]'
+        c << 'log_dir = {0}'.format(self.SHARED_VOLUME)
+        c << 'use_stderr = False'
+        c << '[message]'
+        c << 'write_disk = True'
+        c << 'write_dir = {0}/data/bgp/'.format(self.SHARED_VOLUME)
+        c << 'format = json'
+
+        if self.peers:
+            info = next(iter(self.peers.values()))
+            remote_as = info['remote_as']
+            neigh_addr = info['neigh_addr'].split('/')[0]
+            local_as = info['local_as'] or self.asn
+            local_addr = info['local_addr'].split('/')[0]
+            c << '[bgp]'
+            c << 'afi_safi = ipv4, ipv6, vpnv4, vpnv6, flowspec, evpn'
+            c << 'remote_as = {0}'.format(remote_as)
+            c << 'remote_addr = {0}'.format(neigh_addr)
+            c << 'local_as = {0}'.format(local_as)
+            c << 'local_addr = {0}'.format(local_addr)
+
+        with open('{0}/yabgp.ini'.format(self.config_dir), 'w') as f:
+            print(colors.yellow('[{0}\'s new yabgp.ini]'.format(self.name)))
+            print(colors.yellow(indent(str(c))))
+            f.writelines(str(c))
+
+    def reload_config(self):
+        if self.peers == 0:
+            return
+
+        def _reload():
+            def _is_running():
+                ps = self.local('ps -ef', capture=True)
+                running = False
+                for line in ps.split('\n'):
+                    if 'yabgp_helper' in line:
+                        running = True
+                return running
+
+            if _is_running():
+                self.local('/usr/bin/pkill -9 python')
+
+            self._start_yabgp()
+            time.sleep(self.WAIT_FOR_BOOT)
+            if not _is_running():
+                raise RuntimeError()
+
+        try_several_times(_reload)
+
+    def _curl_send_update(self, path, peer):
+        c = CmdBuffer(' ')
+        c << "curl -X POST"
+        c << "-u admin:admin"
+        c << "-H 'Content-Type: application/json'"
+        c << "http://localhost:8801/v1/peer/{0}/send/update".format(peer)
+        c << "-d '{0}'".format(json.dumps(path))
+        return json.loads(self.local(str(c), capture=True))
+
+    def _construct_ip_unicast_update(self, rf, prefix, nexthop):
+        # YABGP v0.4.0
+        #
+        # IPv4 Unicast:
+        # curl -X POST \
+        #      -u admin:admin \
+        #      -H 'Content-Type: application/json' \
+        #      http://localhost:8801/v1/peer/172.17.0.2/send/update -d '{
+        #     "attr": {
+        #         "1": 0,
+        #         "2": [
+        #             [
+        #                 2,
+        #                 [
+        #                     1,
+        #                     2,
+        #                     3
+        #                 ]
+        #             ]
+        #         ],
+        #         "3": "192.0.2.1",
+        #         "5": 500
+        #     },
+        #     "nlri": [
+        #         "172.20.1.0/24",
+        #         "172.20.2.0/24"
+        #     ]
+        # }'
+        #
+        # IPv6 Unicast:
+        # curl -X POST \
+        #      -u admin:admin \
+        #      -H 'Content-Type: application/json' \
+        #      http://localhost:8801/v1/peer/172.17.0.2/send/update -d '{
+        #     "attr": {
+        #         "1": 0,
+        #         "2": [
+        #             [
+        #                 2,
+        #                 [
+        #                     65502
+        #                 ]
+        #             ]
+        #         ],
+        #         "4": 0,
+        #         "14": {
+        #             "afi_safi": [
+        #                 2,
+        #                 1
+        #             ],
+        #             "linklocal_nexthop": "fe80::c002:bff:fe7e:0",
+        #             "nexthop": "2001:db8::2",
+        #             "nlri": [
+        #                 "::2001:db8:2:2/64",
+        #                 "::2001:db8:2:1/64",
+        #                 "::2001:db8:2:0/64"
+        #             ]
+        #         }
+        #     }
+        # }'
+        if rf == 'ipv4':
+            return {
+                "attr": {
+                    "3": nexthop,
+                },
+                "nlri": [prefix],
+            }
+        elif rf == 'ipv6':
+            return {
+                "attr": {
+                    "14": {  # MP_REACH_NLRI
+                        "afi_safi": [2, 1],
+                        "nexthop": nexthop,
+                        "nlri": [prefix],
+                    },
+                },
+            }
+        else:
+            raise ValueError(
+                'invalid address family for ipv4/ipv6 unicast: %s' % rf)
+
+    def _construct_ip_unicast_withdraw(self, rf, prefix):
+        # YABGP v0.4.0
+        #
+        # IPv4 Unicast:
+        # curl -X POST \
+        #      -u admin:admin \
+        #      -H 'Content-Type: application/json' \
+        #      http://localhost:8801/v1/peer/172.17.0.2/send/update -d '{
+        #     "withdraw": [
+        #         "172.20.1.0/24",
+        #         "172.20.2.0/24"
+        #     ]
+        # }'
+        #
+        # IPv6 Unicast:
+        # curl -X POST \
+        #      -u admin:admin \
+        #      -H 'Content-Type: application/json' \
+        #      http://localhost:8801/v1/peer/172.17.0.2/send/update -d '{
+        #     "attr": {
+        #         "15": {
+        #             "afi_safi": [
+        #                 2,
+        #                 1
+        #             ],
+        #             "withdraw": [
+        #                 "::2001:db8:2:2/64",
+        #                 "::2001:db8:2:1/64",
+        #                 "::2001:db8:2:0/64"
+        #             ]
+        #         }
+        #     }
+        # }'
+        if rf == 'ipv4':
+            return {
+                "withdraw": [prefix],
+            }
+        elif rf == 'ipv6':
+            return {
+                "attr": {
+                    "15": {  # MP_UNREACH_NLRI
+                        "afi_safi": [2, 1],
+                        "withdraw": [prefix],
+                    },
+                },
+            }
+        else:
+            raise ValueError(
+                'invalid address family for ipv4/ipv6 unicast: %s' % rf)
+
+    def _construct_flowspec_match(self, matchs):
+        assert isinstance(matchs, (tuple, list))
+        ret = {}
+        for m in matchs:
+            # m = "source-port '!=2 !=22&!=222'"
+            # typ = "source-port"
+            # args = "'!=2 !=22&!=222'"
+            typ, args = m.split(' ', 1)
+            # t = 6
+            t = FLOWSPEC_NAME_TO_TYPE.get(typ, None)
+            if t is None:
+                raise ValueError('invalid flowspec match type: %s' % typ)
+            # args = "!=2|!=22&!=222"
+            args = args.strip("'").strip('"').replace(' ', '|')
+            ret[t] = args
+        return ret
+
+    def _construct_flowspec_update(self, rf, matchs, thens):
+        # YABGP v0.4.0
+        #
+        # curl -X POST \
+        #      -u admin:admin \
+        #      -H 'Content-Type: application/json' \
+        #      http://localhost:8801/v1/peer/172.17.0.2/send/update -d '{
+        #     "attr": {
+        #         "1": 0,
+        #         "14": {
+        #             "afi_safi": [
+        #                 1,
+        #                 133
+        #             ],
+        #             "nexthop": "",
+        #             "nlri": [
+        #                 {
+        #                     "1": "10.0.0.0/24"
+        #                 }
+        #             ]
+        #         },
+        #         "16": [
+        #             "traffic-rate:0:0"
+        #         ],
+        #         "2": [],
+        #         "5": 100
+        #     }
+        # }'
+        #
+        # Format of "thens":
+        #   "traffic-rate:<AS>:<rate>"
+        #   "traffic-marking-dscp:<int value>"
+        #   "redirect-nexthop:<int value>"
+        #   "redirect-vrf:<RT>"
+        thens = thens or []
+        if rf == 'ipv4-flowspec':
+            afi_safi = [1, 133]
+        else:
+            raise ValueError('invalid address family for flowspec: %s' % rf)
+
+        return {
+            "attr": {
+                "14": {  # MP_REACH_NLRI
+                    "afi_safi": afi_safi,
+                    "nexthop": "",
+                    "nlri": [self._construct_flowspec_match(matchs)]
+                },
+                "16": thens,  # EXTENDED COMMUNITIES
+            },
+        }
+
+    def _construct_flowspec_withdraw(self, rf, matchs):
+        # curl -X POST \
+        #      -u admin:admin \
+        #      -H 'Content-Type: application/json' \
+        #      http://localhost:8801/v1/peer/172.17.0.2/send/update -d '{
+        #     "attr": {
+        #         "15": {
+        #             "afi_safi": [
+        #                 1,
+        #                 133
+        #             ],
+        #             "withdraw": [
+        #                 {
+        #                     "1": "192.88.2.3/24",
+        #                     "2": "192.89.1.3/24"
+        #                 },
+        #                 {
+        #                     "1": "192.88.4.3/24",
+        #                     "2": "192.89.2.3/24"
+        #                 }
+        #             ]
+        #         }
+        #     }
+        # }'
+        if rf == 'ipv4-flowspec':
+            afi_safi = [1, 133]
+        else:
+            raise ValueError('invalid address family for flowspec: %s' % rf)
+
+        return {
+            "attr": {
+                "15": {  # MP_UNREACH_NLRI
+                    "afi_safi": afi_safi,
+                    "withdraw": [self._construct_flowspec_match(matchs)],
+                },
+            },
+        }
+
+    def _update_path_attributes(self, path, aspath=None, med=None,
+                                local_pref=None):
+        # ORIGIN: Currently support only IGP(0)
+        path['attr']['1'] = 0
+        # AS_PATH: Currently support only AS_SEQUENCE(2)
+        if aspath is None:
+            path['attr']['2'] = []
+        else:
+            path['attr']['2'] = [[2, aspath]]
+        # MED
+        if med is not None:
+            path['attr']['4'] = med
+        # LOCAL_PREF
+        if local_pref is not None:
+            path['attr']['5'] = local_pref
+        # TODO:
+        # Support COMMUNITY and EXTENDED COMMUNITIES
+
+        return path
+
+    def add_route(self, route, rf='ipv4', attribute=None, aspath=None,
+                  community=None, med=None, extendedcommunity=None,
+                  nexthop=None, matchs=None, thens=None,
+                  local_pref=None, identifier=None, reload_config=True):
+        self.routes.setdefault(route, [])
+
+        for info in self.peers.values():
+            peer = info['neigh_addr'].split('/')[0]
+
+            if rf in ['ipv4', 'ipv6']:
+                nexthop = nexthop or info['local_addr'].split('/')[0]
+                path = self._construct_ip_unicast_update(
+                    rf, route, nexthop)
+            # TODO:
+            # Support "evpn" address family
+            elif rf in ['ipv4-flowspec', 'ipv6-flowspec']:
+                path = self._construct_flowspec_update(
+                    rf, matchs, thens)
+            else:
+                raise ValueError('unsupported address family: %s' % rf)
+
+            self._update_path_attributes(
+                path, aspath=aspath, med=med, local_pref=local_pref)
+
+            self._curl_send_update(path, peer)
+
+        self.routes[route].append({
+            'prefix': route,
+            'rf': rf,
+            'attr': attribute,
+            'next-hop': nexthop,
+            'as-path': aspath,
+            'community': community,
+            'med': med,
+            'local-pref': local_pref,
+            'extended-community': extendedcommunity,
+            'identifier': identifier,
+            'matchs': matchs,
+            'thens': thens,
+        })
+
+    def del_route(self, route, identifier=None, reload_config=True):
+        new_paths = []
+        withdraw = None
+        for p in self.routes.get(route, []):
+            if p['identifier'] != identifier:
+                new_paths.append(p)
+            else:
+                withdraw = p
+
+        if not withdraw:
+            return
+        rf = withdraw['rf']
+
+        for info in self.peers.values():
+            peer = info['neigh_addr'].split('/')[0]
+
+            if rf in ['ipv4', 'ipv6']:
+                r = self._construct_ip_unicast_withdraw(rf, route)
+            elif rf == 'ipv4-flowspec':
+                # NOTE: "ipv6-flowspec" does not seem to be supported with
+                # YABGP v0.4.0
+                matchs = withdraw['matchs']
+                r = self._construct_flowspec_withdraw(rf, matchs)
+            else:
+                raise ValueError('unsupported address family: %s' % rf)
+
+            self._curl_send_update(r, peer)
+
+        self.routes[route] = new_paths
+
+    def _get_adj_rib(self, peer, in_out='in'):
+        peer_addr = self.peer_name(peer)
+        c = CmdBuffer(' ')
+        c << "curl -X GET"
+        c << "-u admin:admin"
+        c << "-H 'Content-Type: application/json'"
+        c << "http://localhost:8801/v1-ext/peer/{0}/adj-rib-{1}".format(
+            peer_addr, in_out)
+        return json.loads(self.local(str(c), capture=True))
+
+    def get_adj_rib_in(self, peer, rf='ipv4'):
+        # "rf" should be either of;
+        #   ipv4, ipv6, vpnv4, vpnv6, flowspec, evpn
+        # The same as supported "afi_safi" in yabgp.ini
+        ribs = self._get_adj_rib(peer, 'in')
+        return ribs.get(rf, {})
+
+    def get_adj_rib_out(self, peer, rf='ipv4'):
+        # "rf" should be either of;
+        #   ipv4, ipv6, vpnv4, vpnv6, flowspec, evpn
+        # The same as supported "afi_safi" in yabgp.ini
+        ribs = self._get_adj_rib(peer, 'out')
+        return ribs.get(rf, {})
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 = {
+#     <peer.factory.peer_addr>: {
+#         <afi_safi>: {
+#             <nlri>: <msg>,
+#             ...
+#         },
+#         ...
+#     },
+#     ...
+# }
+ADJ_RIB_IN = {}
+
+
+@blueprint.route('/peer/<peer_ip>/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 = {
+#     <peer_ip>: {
+#         <afi_safi>: {
+#             <nlri>: <msg>,
+#             ...
+#         },
+#         ...
+#     },
+#     ...
+# }
+ADJ_RIB_OUT = {}
+
+
+@blueprint.route('/peer/<peer_ip>/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())