| # SPDX-License-Identifier: GPL-2.0 |
| |
| import errno |
| import json as _json |
| import random |
| import re |
| import socket |
| import subprocess |
| import time |
| |
| |
| class CmdExitFailure(Exception): |
| pass |
| |
| |
| class cmd: |
| def __init__(self, comm, shell=True, fail=True, ns=None, background=False, host=None, timeout=5): |
| if ns: |
| comm = f'ip netns exec {ns} ' + comm |
| |
| self.stdout = None |
| self.stderr = None |
| self.ret = None |
| |
| self.comm = comm |
| if host: |
| self.proc = host.cmd(comm) |
| else: |
| self.proc = subprocess.Popen(comm, shell=shell, stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE) |
| if not background: |
| self.process(terminate=False, fail=fail, timeout=timeout) |
| |
| def process(self, terminate=True, fail=None, timeout=5): |
| if fail is None: |
| fail = not terminate |
| |
| if terminate: |
| self.proc.terminate() |
| stdout, stderr = self.proc.communicate(timeout) |
| self.stdout = stdout.decode("utf-8") |
| self.stderr = stderr.decode("utf-8") |
| self.proc.stdout.close() |
| self.proc.stderr.close() |
| self.ret = self.proc.returncode |
| |
| if self.proc.returncode != 0 and fail: |
| if len(stderr) > 0 and stderr[-1] == "\n": |
| stderr = stderr[:-1] |
| raise CmdExitFailure("Command failed: %s\nSTDOUT: %s\nSTDERR: %s" % |
| (self.proc.args, stdout, stderr)) |
| |
| |
| class bkg(cmd): |
| def __init__(self, comm, shell=True, fail=None, ns=None, host=None, |
| exit_wait=False): |
| super().__init__(comm, background=True, |
| shell=shell, fail=fail, ns=ns, host=host) |
| self.terminate = not exit_wait |
| self.check_fail = fail |
| |
| def __enter__(self): |
| return self |
| |
| def __exit__(self, ex_type, ex_value, ex_tb): |
| return self.process(terminate=self.terminate, fail=self.check_fail) |
| |
| |
| global_defer_queue = [] |
| |
| |
| class defer: |
| def __init__(self, func, *args, **kwargs): |
| global global_defer_queue |
| |
| if not callable(func): |
| raise Exception("defer created with un-callable object, did you call the function instead of passing its name?") |
| |
| self.func = func |
| self.args = args |
| self.kwargs = kwargs |
| |
| self._queue = global_defer_queue |
| self._queue.append(self) |
| |
| def __enter__(self): |
| return self |
| |
| def __exit__(self, ex_type, ex_value, ex_tb): |
| return self.exec() |
| |
| def exec_only(self): |
| self.func(*self.args, **self.kwargs) |
| |
| def cancel(self): |
| self._queue.remove(self) |
| |
| def exec(self): |
| self.cancel() |
| self.exec_only() |
| |
| |
| def tool(name, args, json=None, ns=None, host=None): |
| cmd_str = name + ' ' |
| if json: |
| cmd_str += '--json ' |
| cmd_str += args |
| cmd_obj = cmd(cmd_str, ns=ns, host=host) |
| if json: |
| return _json.loads(cmd_obj.stdout) |
| return cmd_obj |
| |
| |
| def ip(args, json=None, ns=None, host=None): |
| if ns: |
| args = f'-netns {ns} ' + args |
| return tool('ip', args, json=json, host=host) |
| |
| |
| def ethtool(args, json=None, ns=None, host=None): |
| return tool('ethtool', args, json=json, ns=ns, host=host) |
| |
| |
| def rand_port(): |
| """ |
| Get a random unprivileged port, try to make sure it's not already used. |
| """ |
| for _ in range(1000): |
| port = random.randint(10000, 65535) |
| try: |
| with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s: |
| s.bind(("", port)) |
| return port |
| except OSError as e: |
| if e.errno != errno.EADDRINUSE: |
| raise |
| raise Exception("Can't find any free unprivileged port") |
| |
| |
| def wait_port_listen(port, proto="tcp", ns=None, host=None, sleep=0.005, deadline=5): |
| end = time.monotonic() + deadline |
| |
| pattern = f":{port:04X} .* " |
| if proto == "tcp": # for tcp protocol additionally check the socket state |
| pattern += "0A" |
| pattern = re.compile(pattern) |
| |
| while True: |
| data = cmd(f'cat /proc/net/{proto}*', ns=ns, host=host, shell=True).stdout |
| for row in data.split("\n"): |
| if pattern.search(row): |
| return |
| if time.monotonic() > end: |
| raise Exception("Waiting for port listen timed out") |
| time.sleep(sleep) |
