1 # SPDX-License-Identifier: GPL-2.0
2
3 import errno
4 import json as _json
5 import random
6 import re
7 import socket
8 import subprocess
9 import time
10
11
12 class CmdExitFailure(Exception):
13 pass
14
15
16 class cmd:
17 def __init__(self, comm, shell=True, fail=True, ns=None, background=False, host=None, timeout=5):
18 if ns:
19 comm = f'ip netns exec {ns} ' + comm
20
21 self.stdout = None
22 self.stderr = None
23 self.ret = None
24
25 self.comm = comm
26 if host:
27 self.proc = host.cmd(comm)
28 else:
29 self.proc = subprocess.Popen(comm, shell=shell, stdout=subprocess.PIPE,
30 stderr=subprocess.PIPE)
31 if not background:
32 self.process(terminate=False, fail=fail, timeout=timeout)
33
34 def process(self, terminate=True, fail=None, timeout=5):
35 if fail is None:
36 fail = not terminate
37
38 if terminate:
39 self.proc.terminate()
40 stdout, stderr = self.proc.communicate(timeout)
41 self.stdout = stdout.decode("utf-8")
42 self.stderr = stderr.decode("utf-8")
43 self.proc.stdout.close()
44 self.proc.stderr.close()
45 self.ret = self.proc.returncode
46
47 if self.proc.returncode != 0 and fail:
48 if len(stderr) > 0 and stderr[-1] == "\n":
49 stderr = stderr[:-1]
50 raise CmdExitFailure("Command failed: %s\nSTDOUT: %s\nSTDERR: %s" %
51 (self.proc.args, stdout, stderr))
52
53
54 class bkg(cmd):
55 def __init__(self, comm, shell=True, fail=None, ns=None, host=None,
56 exit_wait=False):
57 super().__init__(comm, background=True,
58 shell=shell, fail=fail, ns=ns, host=host)
59 self.terminate = not exit_wait
60 self.check_fail = fail
61
62 def __enter__(self):
63 return self
64
65 def __exit__(self, ex_type, ex_value, ex_tb):
66 return self.process(terminate=self.terminate, fail=self.check_fail)
67
68
69 global_defer_queue = []
70
71
72 class defer:
73 def __init__(self, func, *args, **kwargs):
74 global global_defer_queue
75
76 if not callable(func):
77 raise Exception("defer created with un-callable object, did you call the function instead of passing its name?")
78
79 self.func = func
80 self.args = args
81 self.kwargs = kwargs
82
83 self._queue = global_defer_queue
84 self._queue.append(self)
85
86 def __enter__(self):
87 return self
88
89 def __exit__(self, ex_type, ex_value, ex_tb):
90 return self.exec()
91
92 def exec_only(self):
93 self.func(*self.args, **self.kwargs)
94
95 def cancel(self):
96 self._queue.remove(self)
97
98 def exec(self):
99 self.cancel()
100 self.exec_only()
101
102
103 def tool(name, args, json=None, ns=None, host=None):
104 cmd_str = name + ' '
105 if json:
106 cmd_str += '--json '
107 cmd_str += args
108 cmd_obj = cmd(cmd_str, ns=ns, host=host)
109 if json:
110 return _json.loads(cmd_obj.stdout)
111 return cmd_obj
112
113
114 def ip(args, json=None, ns=None, host=None):
115 if ns:
116 args = f'-netns {ns} ' + args
117 return tool('ip', args, json=json, host=host)
118
119
120 def ethtool(args, json=None, ns=None, host=None):
121 return tool('ethtool', args, json=json, ns=ns, host=host)
122
123
124 def rand_port():
125 """
126 Get a random unprivileged port, try to make sure it's not already used.
127 """
128 for _ in range(1000):
129 port = random.randint(10000, 65535)
130 try:
131 with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
132 s.bind(("", port))
133 return port
134 except OSError as e:
135 if e.errno != errno.EADDRINUSE:
136 raise
137 raise Exception("Can't find any free unprivileged port")
138
139
140 def wait_port_listen(port, proto="tcp", ns=None, host=None, sleep=0.005, deadline=5):
141 end = time.monotonic() + deadline
142
143 pattern = f":{port:04X} .* "
144 if proto == "tcp": # for tcp protocol additionally check the socket state
145 pattern += "0A"
146 pattern = re.compile(pattern)
147
148 while True:
149 data = cmd(f'cat /proc/net/{proto}*', ns=ns, host=host, shell=True).stdout
150 for row in data.split("\n"):
151 if pattern.search(row):
152 return
153 if time.monotonic() > end:
154 raise Exception("Waiting for port listen timed out")
155 time.sleep(sleep)
Linux® is a registered trademark of Linus Torvalds in the United States and other countries.
TOMOYO® is a registered trademark of NTT DATA CORPORATION.