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) 156