1# SPDX-License-Identifier: GPL-2.0
2
3import json as _json
4import random
5import re
6import subprocess
7import time
8
9
10class cmd:
11    def __init__(self, comm, shell=True, fail=True, ns=None, background=False, host=None, timeout=5):
12        if ns:
13            comm = f'ip netns exec {ns} ' + comm
14
15        self.stdout = None
16        self.stderr = None
17        self.ret = None
18
19        self.comm = comm
20        if host:
21            self.proc = host.cmd(comm)
22        else:
23            self.proc = subprocess.Popen(comm, shell=shell, stdout=subprocess.PIPE,
24                                         stderr=subprocess.PIPE)
25        if not background:
26            self.process(terminate=False, fail=fail, timeout=timeout)
27
28    def process(self, terminate=True, fail=None, timeout=5):
29        if fail is None:
30            fail = not terminate
31
32        if terminate:
33            self.proc.terminate()
34        stdout, stderr = self.proc.communicate(timeout)
35        self.stdout = stdout.decode("utf-8")
36        self.stderr = stderr.decode("utf-8")
37        self.proc.stdout.close()
38        self.proc.stderr.close()
39        self.ret = self.proc.returncode
40
41        if self.proc.returncode != 0 and fail:
42            if len(stderr) > 0 and stderr[-1] == "\n":
43                stderr = stderr[:-1]
44            raise Exception("Command failed: %s\nSTDOUT: %s\nSTDERR: %s" %
45                            (self.proc.args, stdout, stderr))
46
47
48class bkg(cmd):
49    def __init__(self, comm, shell=True, fail=None, ns=None, host=None,
50                 exit_wait=False):
51        super().__init__(comm, background=True,
52                         shell=shell, fail=fail, ns=ns, host=host)
53        self.terminate = not exit_wait
54        self.check_fail = fail
55
56    def __enter__(self):
57        return self
58
59    def __exit__(self, ex_type, ex_value, ex_tb):
60        return self.process(terminate=self.terminate, fail=self.check_fail)
61
62
63def tool(name, args, json=None, ns=None, host=None):
64    cmd_str = name + ' '
65    if json:
66        cmd_str += '--json '
67    cmd_str += args
68    cmd_obj = cmd(cmd_str, ns=ns, host=host)
69    if json:
70        return _json.loads(cmd_obj.stdout)
71    return cmd_obj
72
73
74def ip(args, json=None, ns=None, host=None):
75    if ns:
76        args = f'-netns {ns} ' + args
77    return tool('ip', args, json=json, host=host)
78
79
80def rand_port():
81    """
82    Get unprivileged port, for now just random, one day we may decide to check if used.
83    """
84    return random.randint(10000, 65535)
85
86
87def wait_port_listen(port, proto="tcp", ns=None, host=None, sleep=0.005, deadline=5):
88    end = time.monotonic() + deadline
89
90    pattern = f":{port:04X} .* "
91    if proto == "tcp": # for tcp protocol additionally check the socket state
92        pattern += "0A"
93    pattern = re.compile(pattern)
94
95    while True:
96        data = cmd(f'cat /proc/net/{proto}*', ns=ns, host=host, shell=True).stdout
97        for row in data.split("\n"):
98            if pattern.search(row):
99                return
100        if time.monotonic() > end:
101            raise Exception("Waiting for port listen timed out")
102        time.sleep(sleep)
103