iproute: Merge pull request #1247 from svinota/1203-iproute-save

iproute: route_dump() and route_load() functionality

Bug-Url: #1247
Bug-Url: #1203

pyroute2/iproute/linux.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
+import io
 import logging
 import os
+import struct
 import time
 import warnings
 from functools import partial
@@ -91,9 +93,10 @@
 from pyroute2.requests.rule import RuleFieldFilter, RuleIPRouteFilter
 from pyroute2.requests.tc import TcIPRouteFilter, TcRequestFilter
 
-from .parsers import default_routes
+from .parsers import default_routes, export_routes
 
 DEFAULT_TABLE = 254
+IPROUTE2_DUMP_MAGIC = 0x45311224
 log = logging.getLogger(__name__)
 
 
@@ -405,6 +408,101 @@ async def probe(self, command, **kwarg):
         await request.send()
         return [x async for x in request.response()]
 
+    # 8<---------------------------------------------------------------
+    #
+    # Binary streams methods
+    #
+    async def route_dump(self, fd, family=AF_UNSPEC, fmt='iproute2'):
+        '''Save routes as a binary dump into a file object.
+
+        fd -- an open file object, must support `write()`
+        family -- AF_UNSPEC, AF_INET, etc. -- filter routes by family
+        fmt -- dump format, "iproute2" (default) or "raw"
+
+        The binary dump is just a set of unparsed netlink messages.
+        The `iproute2` prepends the dump with a magic uint32, so
+        `IPRoute` does the same for compatibility. If you want a raw
+        dump without any additional magic data, use `fmt="raw"`.
+
+        This routine neither close the file object, nor uses `seek()`
+        to rewind, it's up to the user.
+        '''
+
+        if fmt == 'iproute2':
+            fd.write(struct.pack('I', IPROUTE2_DUMP_MAGIC))
+        elif fmt != 'raw':
+            raise TypeError('dump format not supported')
+        msg = rtmsg()
+        msg['family'] = family
+        request = NetlinkRequest(
+            self,
+            msg,
+            msg_type=RTM_GETROUTE,
+            msg_flags=NLM_F_DUMP | NLM_F_REQUEST,
+            parser=export_routes(fd),
+        )
+        await request.send()
+        return [x async for x in request.response()]
+
+    async def route_dumps(self, family=AF_UNSPEC, fmt='iproute2'):
+        '''Save routes and returns as a `bytes` object.
+
+        The same as `.route_dump()`, but returns `bytes`.
+        '''
+        fd = io.BytesIO()
+        await self.route_dump(fd, family, fmt)
+        return fd.getvalue()
+
+    async def route_load(self, fd, fmt='iproute2'):
+        '''Load routes from a binary dump.
+
+        fd -- an open file object, must support `read()`
+        fmt -- dump format, "iproute2" (default) or "raw"
+
+        The current version parses the dump and loads routes one
+        by one. This behavior will be changed in the future to
+        optimize the performance, but the result will be the same.
+
+        If `fmt == "iproute2"`, then the loader checks the magic iproute2
+        prefix in the dump. Otherwise it parses the data from byte 0.
+        '''
+        if fmt == 'iproute2':
+            if (
+                not struct.unpack('I', fd.read(struct.calcsize('I')))[0]
+                == IPROUTE2_DUMP_MAGIC
+            ):
+                raise TypeError('wrong dump magic')
+        elif fmt != 'raw':
+            raise TypeError('dump format not supported')
+        ret = []
+        for msg in self.marshal.parse(fd.read()):
+            request = NetlinkRequest(
+                self,
+                msg,
+                command='replace',
+                command_map={'replace': (RTM_NEWROUTE, 'replace')},
+            )
+            await request.send()
+            ret.extend(
+                [
+                    x['header']['error'] is None
+                    async for x in request.response()
+                ]
+            )
+        if not all(ret):
+            raise NetlinkError('error loading route dump')
+        return []
+
+    async def route_loads(self, data, fmt='iproute2'):
+        '''Load routes from a `bytes` object.
+
+        Like `.route_load()`, but accepts `bytes` instead of an file file.
+        '''
+        fd = io.BytesIO()
+        fd.write(data)
+        fd.seek(0)
+        return await self.route_load(fd, fmt)
+
     # 8<---------------------------------------------------------------
     #
     # Listing methods
@@ -2458,6 +2556,10 @@ def __getattr__(self, name):
             'flush_rules',
             'flush_routes',
             'get_netnsid',
+            'route_dump',
+            'route_dumps',
+            'route_load',
+            'route_loads',
         ]
         async_dump_methods = [
             'dump',

pyroute2/iproute/parsers.py

@@ -1,9 +1,49 @@
 import struct
+from functools import partial
 
 from pyroute2.netlink import NLMSG_DONE, nlmsg
+from pyroute2.netlink.exceptions import NetlinkError
+from pyroute2.netlink.rtnl import RTM_NEWROUTE
 from pyroute2.netlink.rtnl.rtmsg import rtmsg
 
 
+def get_header(data, offset):
+    # get message header
+    header = dict(
+        zip(
+            ('length', 'type', 'flags', 'sequence_number'),
+            struct.unpack_from('IHHI', data, offset),
+        )
+    )
+    header['error'] = None
+    return header
+
+
+def msg_done(header):
+    msg = nlmsg()
+    msg['header'] = header
+    msg.length = msg['header']['length']
+    return msg
+
+
+def _export_routes(fd, data, offset, length):
+    '''Export RTM_NEWROUTE messages binary data.
+
+    Otherwise return NLMSG_DONE.
+    '''
+    header = get_header(data, offset)
+    if header['type'] == NLMSG_DONE:
+        return msg_done(header)
+    elif header['type'] == RTM_NEWROUTE:
+        fd.write(data[offset : offset + length])
+        return
+    raise NetlinkError()
+
+
+def export_routes(fd):
+    return partial(_export_routes, fd)
+
+
 def default_routes(data, offset, length):
     '''
     Only for RTM_NEWROUTE.
@@ -14,19 +54,9 @@ def default_routes(data, offset, length):
     * nlmsg() -- NLMSG_DONE
     * None for any other messages
     '''
-    # get message header
-    header = dict(
-        zip(
-            ('length', 'type', 'flags', 'sequence_number'),
-            struct.unpack_from('IHHI', data, offset),
-        )
-    )
-    header['error'] = None
+    header = get_header(data, offset)
     if header['type'] == NLMSG_DONE:
-        msg = nlmsg()
-        msg['header'] = header
-        msg.length = msg['header']['length']
-        return msg
+        return msg_done(header)
 
     # skip to NLA: offset + nlmsg header + rtmsg data
     cursor = offset + 28

tests/test_core/conftest.py

@@ -1,16 +1,22 @@
 import errno
+import os
 
 import pytest
 import pytest_asyncio
 from pr2test.plan9 import AsyncPlan9Context
 
-from pyroute2 import AsyncIPRoute, IPRoute, NetlinkError
+from pyroute2 import AsyncIPRoute, IPRoute, NetlinkError, netns
 from pyroute2.common import uifname
 
 
 class AsyncIPRouteContext(AsyncIPRoute):
     def __init__(self, *argv, **kwarg):
+        self.remove_netns_on_exit = False
         self.registry_ifname = set()
+        if kwarg.get('netns') is True:
+            kwarg['netns'] = uifname()
+            kwarg['flags'] = os.O_CREAT
+            self.remove_netns_on_exit = True
         super().__init__(*argv, **kwarg)
 
     def register_temporary_ifname(self, ifname=None):
@@ -25,25 +31,40 @@ async def close(self, *argv, **kwarg):
             except NetlinkError as e:
                 if e.code != errno.ENODEV:
                     raise
+        await super().close(*argv, **kwarg)
+        if self.remove_netns_on_exit:
+            netns.remove(self.status['netns'])
 
 
 class SyncIPRouteContext(IPRoute):
     def __init__(self, *argv, **kwarg):
+        self.remove_netns_on_exit = False
         self.registry_ifname = set()
+        if kwarg.get('netns') is True:
+            kwarg['netns'] = uifname()
+            kwarg['flags'] = os.O_CREAT
+            self.remove_netns_on_exit = True
         super().__init__(*argv, **kwarg)
 
     def register_temporary_ifname(self, ifname=None):
         ifname = ifname if ifname is not None else uifname()
         self.registry_ifname.add(ifname)
         return ifname
 
+    def register_temporary_netns(self, netns=None):
+        netns = netns if netns is not None else uifname()
+        self.registry_netns.add(netns)
+
     def close(self, *argv, **kwarg):
         for ifname in self.registry_ifname:
             try:
                 self.link('del', ifname=ifname)
             except NetlinkError as e:
                 if e.code != errno.ENODEV:
                     raise
+        super().close(*argv, **kwarg)
+        if self.remove_netns_on_exit:
+            netns.remove(self.status['netns'])
 
 
 @pytest_asyncio.fixture
@@ -56,11 +77,13 @@ async def p9(request, tmpdir):
 
 @pytest_asyncio.fixture
 async def async_ipr(request, tmpdir):
-    async with AsyncIPRouteContext() as ctx:
+    kwarg = getattr(request, 'param', {})
+    async with AsyncIPRouteContext(**kwarg) as ctx:
         yield ctx
 
 
 @pytest.fixture
 def sync_ipr(request, tmpdir):
-    with SyncIPRouteContext() as ctx:
+    kwarg = getattr(request, 'param', {})
+    with SyncIPRouteContext(**kwarg) as ctx:
         yield ctx

tests/test_core/pr2test/tools.py

@@ -52,18 +52,31 @@ def wait_for_ip_object(cmd, filters, timeout, retry):
     return found
 
 
-def address_exists(address, ifname=None, timeout=1, retry=0.2):
+def address_exists(address, ifname=None, netns=None, timeout=1, retry=0.2):
+    ns = [] if netns is None else ['ip', 'netns', 'exec', netns]
     filters = [ip_object_filter(query='.addr_info.local', value=address)]
     ifspec = ['dev', ifname] if ifname is not None else []
     return wait_for_ip_object(
-        ['ip', '-json', 'addr', 'show'] + ifspec, filters, timeout, retry
+        ns + ['ip', '-json', 'addr', 'show'] + ifspec, filters, timeout, retry
     )
 
 
 def interface_exists(ifname, netns=None, timeout=1, retry=0.2):
+    ns = [] if netns is None else ['ip', 'netns', 'exec', netns]
     filters = [ip_object_filter(query='.ifname', value=ifname)]
     return wait_for_ip_object(
-        ['ip', '-json', 'link', 'show'], filters, timeout, retry
+        ns + ['ip', '-json', 'link', 'show'], filters, timeout, retry
+    )
+
+
+def route_exists(dst, table='main', netns=None, timeout=1, retry=0.2):
+    ns = [] if netns is None else ['ip', 'netns', 'exec', netns]
+    filters = [ip_object_filter(query='.dst', value=dst)]
+    return wait_for_ip_object(
+        ns + ['ip', '-json', 'route', 'show', 'table', str(table)],
+        filters,
+        timeout,
+        retry,
     )