1 files changed, 109 insertions, 0 deletions
diff --git a/test/lib/quagga.py b/test/lib/quagga.py
index 897648ae..e7b9f9bf 100644
--- a/test/lib/quagga.py
+++ b/test/lib/quagga.py
@@ -21,6 +21,7 @@ import netaddr
 
 from lib.base import (
     BGPContainer,
+    OSPFContainer,
     CmdBuffer,
     BGP_FSM_IDLE,
     BGP_FSM_ACTIVE,
@@ -282,3 +283,111 @@ class RawQuaggaBGPContainer(QuaggaBGPContainer):
             print colors.yellow('[{0}\'s new bgpd.conf]'.format(self.name))
             print colors.yellow(indent(self.config))
             f.writelines(self.config)
+
+
+class QuaggaOSPFContainer(OSPFContainer):
+    SHARED_VOLUME = '/etc/quagga'
+    ZAPI_V2_IMAGE = 'osrg/quagga'
+    ZAPI_V3_IMAGE = 'osrg/quagga:v1.0'
+
+    def __init__(self, name, image=ZAPI_V2_IMAGE, zapi_verion=2,
+                 zebra_config=None, ospfd_config=None):
+        if zapi_verion != 2:
+            image = self.ZAPI_V3_IMAGE
+        super(QuaggaOSPFContainer, self).__init__(name, image)
+        self.shared_volumes.append((self.config_dir, self.SHARED_VOLUME))
+
+        self.zapi_vserion = zapi_verion
+
+        # Example:
+        # zebra_config = {
+        #     'interfaces': {  # interface settings
+        #         'eth0': [
+        #             'ip address 192.168.0.1/24',
+        #         ],
+        #     },
+        #     'routes': [  # static route settings
+        #         'ip route 172.16.0.0/16 172.16.0.1',
+        #     ],
+        # }
+        self.zebra_config = zebra_config or {}
+
+        # Example:
+        # ospfd_config = {
+        #     'redistribute_types': [
+        #         'connected',
+        #     ],
+        #     'networks': {
+        #         '192.168.1.0/24': '0.0.0.0',  # <network>: <area>
+        #     },
+        # }
+        self.ospfd_config = ospfd_config or {}
+
+    def run(self):
+        super(QuaggaOSPFContainer, self).run()
+        # self.create_config() is called in super(...).run()
+        self._start_zebra()
+        self._start_ospfd()
+        return self.WAIT_FOR_BOOT
+
+    def create_config(self):
+        self._create_config_zebra()
+        self._create_config_ospfd()
+
+    def _create_config_zebra(self):
+        c = CmdBuffer()
+        c << 'hostname zebra'
+        c << 'password zebra'
+        for name, settings in self.zebra_config.get('interfaces', {}).items():
+            c << 'interface {0}'.format(name)
+            for setting in settings:
+                c << str(setting)
+        for route in self.zebra_config.get('routes', []):
+            c << str(route)
+        c << 'log file {0}/zebra.log'.format(self.SHARED_VOLUME)
+        c << 'debug zebra packet'
+        c << 'debug zebra kernel'
+        c << 'debug zebra rib'
+        c << ''
+
+        with open('{0}/zebra.conf'.format(self.config_dir), 'w') as f:
+            print colors.yellow('[{0}\'s new zebra.conf]'.format(self.name))
+            print colors.yellow(indent(str(c)))
+            f.writelines(str(c))
+
+    def _create_config_ospfd(self):
+        c = CmdBuffer()
+        c << 'hostname ospfd'
+        c << 'password zebra'
+        c << 'router ospf'
+        for redistribute in self.ospfd_config.get('redistributes', []):
+            c << ' redistribute {0}'.format(redistribute)
+        for network, area in self.ospfd_config.get('networks', {}).items():
+            self.networks[network] = area  # for superclass
+            c << ' network {0} area {1}'.format(network, area)
+        c << 'log file {0}/ospfd.log'.format(self.SHARED_VOLUME)
+        c << ''
+
+        with open('{0}/ospfd.conf'.format(self.config_dir), 'w') as f:
+            print colors.yellow('[{0}\'s new ospfd.conf]'.format(self.name))
+            print colors.yellow(indent(str(c)))
+            f.writelines(str(c))
+
+    def _start_zebra(self):
+        if self.zapi_vserion == 2:
+            # Do nothing. supervisord will automatically start Zebra daemon.
+            return
+
+        self.local(
+            '/usr/local/sbin/zebra '
+            '-u root -g root -f {0}/zebra.conf'.format(self.SHARED_VOLUME),
+            detach=True)
+
+    def _start_ospfd(self):
+        if self.zapi_vserion == 2:
+            ospfd_cmd = '/usr/lib/quagga/ospfd'
+        else:
+            ospfd_cmd = '/usr/local/sbin/ospfd -u root -g root'
+        self.local(
+            '{0} -f {1}/ospfd.conf'.format(ospfd_cmd, self.SHARED_VOLUME),
+            detach=True)