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