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