1import os
2import signal
3from string import Template
4import subprocess
5import time
6from multiprocessing import Pool
7from functools import cached_property
8from TdcPlugin import TdcPlugin
9
10from tdc_config import *
11
12try:
13    from pyroute2 import netns
14    from pyroute2 import IPRoute
15    netlink = True
16except ImportError:
17    netlink = False
18    print("!!! Consider installing pyroute2 !!!")
19
20class SubPlugin(TdcPlugin):
21    def __init__(self):
22        self.sub_class = 'ns/SubPlugin'
23        super().__init__()
24
25    def pre_suite(self, testcount, testlist):
26        super().pre_suite(testcount, testlist)
27
28    def prepare_test(self, test):
29        if 'skip' in test and test['skip'] == 'yes':
30            return
31
32        if 'nsPlugin' not in test['plugins']:
33            return
34
35        if netlink == True:
36            self._nl_ns_create()
37        else:
38            self._ipr2_ns_create()
39
40        # Make sure the netns is visible in the fs
41        ticks = 20
42        while True:
43            if ticks == 0:
44                raise TimeoutError
45            self._proc_check()
46            try:
47                ns = self.args.NAMES['NS']
48                f = open('/run/netns/{}'.format(ns))
49                f.close()
50                break
51            except:
52                time.sleep(0.1)
53                ticks -= 1
54                continue
55
56    def pre_case(self, test, test_skip):
57        if self.args.verbose:
58            print('{}.pre_case'.format(self.sub_class))
59
60        if test_skip:
61            return
62
63        self.prepare_test(test)
64
65    def post_case(self):
66        if self.args.verbose:
67            print('{}.post_case'.format(self.sub_class))
68
69        if netlink == True:
70            self._nl_ns_destroy()
71        else:
72            self._ipr2_ns_destroy()
73
74    def post_suite(self, index):
75        if self.args.verbose:
76            print('{}.post_suite'.format(self.sub_class))
77
78        # Make sure we don't leak resources
79        cmd = self._replace_keywords("$IP -a netns del")
80
81        if self.args.verbose > 3:
82            print('_exec_cmd:  command "{}"'.format(cmd))
83
84        subprocess.run(cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
85
86    def adjust_command(self, stage, command):
87        super().adjust_command(stage, command)
88        cmdform = 'list'
89        cmdlist = list()
90
91        if self.args.verbose:
92            print('{}.adjust_command'.format(self.sub_class))
93
94        if not isinstance(command, list):
95            cmdform = 'str'
96            cmdlist = command.split()
97        else:
98            cmdlist = command
99        if stage == 'setup' or stage == 'execute' or stage == 'verify' or stage == 'teardown':
100            if self.args.verbose:
101                print('adjust_command:  stage is {}; inserting netns stuff in command [{}] list [{}]'.format(stage, command, cmdlist))
102            cmdlist.insert(0, self.args.NAMES['NS'])
103            cmdlist.insert(0, 'exec')
104            cmdlist.insert(0, 'netns')
105            cmdlist.insert(0, self.args.NAMES['IP'])
106        else:
107            pass
108
109        if cmdform == 'str':
110            command = ' '.join(cmdlist)
111        else:
112            command = cmdlist
113
114        if self.args.verbose:
115            print('adjust_command:  return command [{}]'.format(command))
116        return command
117
118    def _nl_ns_create(self):
119        ns = self.args.NAMES["NS"];
120        dev0 = self.args.NAMES["DEV0"];
121        dev1 = self.args.NAMES["DEV1"];
122        dummy = self.args.NAMES["DUMMY"];
123
124        if self.args.verbose:
125            print('{}._nl_ns_create'.format(self.sub_class))
126
127        netns.create(ns)
128        netns.pushns(newns=ns)
129        with IPRoute() as ip:
130            ip.link('add', ifname=dev1, kind='veth', peer={'ifname': dev0, 'net_ns_fd':'/proc/1/ns/net'})
131            ip.link('add', ifname=dummy, kind='dummy')
132            ticks = 20
133            while True:
134                if ticks == 0:
135                    raise TimeoutError
136                try:
137                    dev1_idx = ip.link_lookup(ifname=dev1)[0]
138                    dummy_idx = ip.link_lookup(ifname=dummy)[0]
139                    ip.link('set', index=dev1_idx, state='up')
140                    ip.link('set', index=dummy_idx, state='up')
141                    break
142                except:
143                    time.sleep(0.1)
144                    ticks -= 1
145                    continue
146        netns.popns()
147
148        with IPRoute() as ip:
149            ticks = 20
150            while True:
151                if ticks == 0:
152                    raise TimeoutError
153                try:
154                    dev0_idx = ip.link_lookup(ifname=dev0)[0]
155                    ip.link('set', index=dev0_idx, state='up')
156                    break
157                except:
158                    time.sleep(0.1)
159                    ticks -= 1
160                    continue
161
162    def _ipr2_ns_create_cmds(self):
163        cmds = []
164
165        ns = self.args.NAMES['NS']
166
167        cmds.append(self._replace_keywords('netns add {}'.format(ns)))
168        cmds.append(self._replace_keywords('link add $DEV1 type veth peer name $DEV0'))
169        cmds.append(self._replace_keywords('link set $DEV1 netns {}'.format(ns)))
170        cmds.append(self._replace_keywords('link add $DUMMY type dummy'.format(ns)))
171        cmds.append(self._replace_keywords('link set $DUMMY netns {}'.format(ns)))
172        cmds.append(self._replace_keywords('netns exec {} $IP link set $DEV1 up'.format(ns)))
173        cmds.append(self._replace_keywords('netns exec {} $IP link set $DUMMY up'.format(ns)))
174        cmds.append(self._replace_keywords('link set $DEV0 up'.format(ns)))
175
176        if self.args.device:
177            cmds.append(self._replace_keywords('link set $DEV2 netns {}'.format(ns)))
178            cmds.append(self._replace_keywords('netns exec {} $IP link set $DEV2 up'.format(ns)))
179
180        return cmds
181
182    def _ipr2_ns_create(self):
183        '''
184        Create the network namespace in which the tests will be run and set up
185        the required network devices for it.
186        '''
187        self._exec_cmd_batched('pre', self._ipr2_ns_create_cmds())
188
189    def _nl_ns_destroy(self):
190        ns = self.args.NAMES['NS']
191        netns.remove(ns)
192
193    def _ipr2_ns_destroy_cmd(self):
194        return self._replace_keywords('netns delete {}'.format(self.args.NAMES['NS']))
195
196    def _ipr2_ns_destroy(self):
197        '''
198        Destroy the network namespace for testing (and any associated network
199        devices as well)
200        '''
201        self._exec_cmd('post', self._ipr2_ns_destroy_cmd())
202
203    @cached_property
204    def _proc(self):
205        ip = self._replace_keywords("$IP -b -")
206        proc = subprocess.Popen(ip,
207            shell=True,
208            stdin=subprocess.PIPE,
209            env=ENVIR)
210
211        return proc
212
213    def _proc_check(self):
214        proc = self._proc
215
216        proc.poll()
217
218        if proc.returncode is not None and proc.returncode != 0:
219            raise RuntimeError("iproute2 exited with an error code")
220
221    def _exec_cmd(self, stage, command):
222        '''
223        Perform any required modifications on an executable command, then run
224        it in a subprocess and return the results.
225        '''
226
227        if self.args.verbose > 3:
228            print('_exec_cmd:  command "{}"'.format(command))
229
230        proc = self._proc
231
232        proc.stdin.write((command + '\n').encode())
233        proc.stdin.flush()
234
235        if self.args.verbose > 3:
236            print('_exec_cmd proc: {}'.format(proc))
237
238        self._proc_check()
239
240    def _exec_cmd_batched(self, stage, commands):
241        for cmd in commands:
242            self._exec_cmd(stage, cmd)
243
244    def _replace_keywords(self, cmd):
245        """
246        For a given executable command, substitute any known
247        variables contained within NAMES with the correct values
248        """
249        tcmd = Template(cmd)
250        subcmd = tcmd.safe_substitute(self.args.NAMES)
251        return subcmd
252