# SPDX-License-Identifier: GPL-2.0

import json as _json
import os
import re
import select
import socket
import subprocess
import time


class CmdInitFailure(Exception):
    """ Command failed to start. Only raised by bkg(). """
    def __init__(self, msg, cmd_obj):
        super().__init__(msg + "\n" + repr(cmd_obj))
        self.cmd = cmd_obj


class CmdExitFailure(Exception):
    """ Command failed (returned non-zero exit code). """
    def __init__(self, msg, cmd_obj):
        super().__init__(msg + "\n" + repr(cmd_obj))
        self.cmd = cmd_obj


class CmdExitZeroFailure(CmdExitFailure):
    """ Command succeeded (returned zero exit code), but expected failure. """


def fd_read_timeout(fd, timeout):
    rlist, _, _ = select.select([fd], [], [], timeout)
    if rlist:
        return os.read(fd, 1024)
    raise TimeoutError("Timeout waiting for fd read")


class cmd:
    """
    Execute a command on local or remote host.

    @shell defaults to false, and class will try to split @comm into a list
    if it's a string with spaces.

    Use bkg() instead to run a command in the background.
    """
    def __init__(self, comm, shell=None, fail=True, expect_fail=False, ns=None,
                 background=False, host=None, timeout=20, ksft_ready=None,
                 ksft_wait=None):
        if ns:
            if hasattr(ns, 'user_ns_path'):
                comm = (f'nsenter --user={ns.user_ns_path} '
                        f'--net={ns.net_ns_path} --setuid=0 --setgid=0 -- '
                        + comm)
            else:
                comm = f'ip netns exec {ns} ' + comm

        self.stdout = None
        self.stderr = None
        self.ret = None
        self.ksft_term_fd = None

        self.host = host
        self.comm = comm

        if host:
            self.proc = host.cmd(comm)
        else:
            # If user doesn't explicitly request shell try to avoid it.
            if shell is None and isinstance(comm, str) and ' ' in comm:
                comm = comm.split()

            # ksft_wait lets us wait for the background process to fully start,
            # we pass an FD to the child process, and wait for it to write back.
            # Similarly term_fd tells child it's time to exit.
            pass_fds = []
            env = os.environ.copy()
            if ksft_wait is not None:
                wait_fd, self.ksft_term_fd = os.pipe()
                pass_fds.append(wait_fd)
                env["KSFT_WAIT_FD"]  = str(wait_fd)
                ksft_ready = True  # ksft_wait implies ready
            if ksft_ready is not None:
                rfd, ready_fd = os.pipe()
                pass_fds.append(ready_fd)
                env["KSFT_READY_FD"] = str(ready_fd)

            self.proc = subprocess.Popen(comm, shell=shell, stdout=subprocess.PIPE,
                                         stderr=subprocess.PIPE, pass_fds=pass_fds,
                                         env=env)
            if ksft_wait is not None:
                os.close(wait_fd)
            if ksft_ready is not None:
                os.close(ready_fd)
                msg = fd_read_timeout(rfd, ksft_wait)
                os.close(rfd)
                if not msg:
                    terminate = self.proc.poll() is None
                    self._process_terminate(terminate=terminate, timeout=1)
                    raise CmdInitFailure("Did not receive ready message", self)
        if not background:
            self.process(terminate=False, fail=fail, expect_fail=expect_fail,
                         timeout=timeout)

    def _process_terminate(self, terminate, timeout):
        if terminate:
            self.proc.terminate()
        stdout, stderr = self.proc.communicate(timeout=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

        return stdout, stderr

    def process(self, terminate=True, fail=None, expect_fail=False, timeout=20):
        if fail is None:
            fail = not terminate

        if self.ksft_term_fd:
            os.write(self.ksft_term_fd, b"1")

        stdout, stderr = self._process_terminate(terminate=terminate,
                                                 timeout=timeout)

        # Fail on unexpected test failure if fail.
        # Fail on unexpected test success if expect_fail.
        # Fail on negative returncode if either:
        # Set by subprocess on crash or signal, this is never expected failure.
        if (self.proc.returncode != 0 and fail or
            (self.proc.returncode < 0 and expect_fail)):
            if len(stderr) > 0 and stderr[-1] == "\n":
                stderr = stderr[:-1]
            raise CmdExitFailure("Command failed", self)
        elif self.proc.returncode == 0 and expect_fail:
            raise CmdExitZeroFailure("Command succeeded (expected fail)", self)


    def __repr__(self):
        def str_fmt(name, s):
            name += ': '
            return (name + s.strip().replace('\n', '\n' + ' ' * len(name)))

        ret = "CMD"
        if self.host:
            ret += "[remote]"
        if self.ret is None:
            ret += f" (unterminated): {self.comm}\n"
        elif self.ret == 0:
            ret += f" (success): {self.comm}\n"
        else:
            ret += f": {self.comm}\n"
            ret += f"  EXIT: {self.ret}\n"
        if self.stdout:
            ret += str_fmt("  STDOUT", self.stdout) + "\n"
        if self.stderr:
            ret += str_fmt("  STDERR", self.stderr) + "\n"
        return ret.strip()


class bkg(cmd):
    """
    Run a command in the background.

    Examples usage:

    Run a command on remote host, and wait for it to finish.
    This is usually paired with wait_port_listen() to make sure
    the command has initialized:

        with bkg("socat ...", exit_wait=True, host=cfg.remote) as nc:
            ...

    Run a command and expect it to let us know that it's ready
    by writing to a special file descriptor passed via KSFT_READY_FD.
    Command will be terminated when we exit the context manager:

        with bkg("my_binary", ksft_wait=5):
    """
    def __init__(self, comm, shell=None, fail=None, expect_fail=None,
                 ns=None, host=None, exit_wait=False, ksft_ready=None,
                 ksft_wait=None):
        super().__init__(comm, background=True,
                         shell=shell, fail=fail, expect_fail=expect_fail,
                         ns=ns, host=host, ksft_ready=ksft_ready,
                         ksft_wait=ksft_wait)
        self.terminate = not exit_wait and not ksft_wait
        self._exit_wait = exit_wait
        self.check_fail = fail
        self.expect_fail = expect_fail

        if shell and self.terminate:
            print("# Warning: combining shell and terminate is risky!")
            print("#          SIGTERM may not reach the child on zsh/ksh!")

    def __enter__(self):
        return self

    def __exit__(self, ex_type, ex_value, ex_tb):
        terminate = self.terminate
        # Force termination on exception, but only if bkg() didn't already exit
        # since forcing termination silences failures with fail=None
        if self.proc.poll() is None:
            terminate = terminate or (self._exit_wait and ex_type is not None)
        return self.process(terminate=terminate, fail=self.check_fail,
                            expect_fail=self.expect_fail)


GLOBAL_DEFER_QUEUE = []
GLOBAL_DEFER_ARMED = False


class defer:
    def __init__(self, func, *args, **kwargs):
        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

        if not GLOBAL_DEFER_ARMED:
            raise Exception("defer queue not armed, did you use defer() outside of a test case?")
        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:
        if name == 'tc':
            cmd_str += '-json '
        else:
            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 bpftool(args, json=None, ns=None, host=None):
    return tool('bpftool', args, json=json, ns=ns, host=host)


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 tc(args, json=None, ns=None, host=None):
    """ Helper to call tc with standard set of optional args. """
    if ns:
        args = f'-netns {ns} ' + args
    return tool('tc', 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 bpftrace(expr, json=None, ns=None, host=None, timeout=None):
    """
    Run bpftrace and return map data (if json=True).
    The output of bpftrace is inconvenient, so the helper converts
    to a dict indexed by map name, e.g.:
     {
       "@":     { ... },
       "@map2": { ... },
     }
    """
    cmd_arr = ['bpftrace']
    # Throw in --quiet if json, otherwise the output has two objects
    if json:
        cmd_arr += ['-f', 'json', '-q']
    if timeout:
        expr += ' interval:s:' + str(timeout) + ' { exit(); }'
        timeout += 20
    cmd_arr += ['-e', expr]
    cmd_obj = cmd(cmd_arr, ns=ns, host=host, shell=False, timeout=timeout)
    if json:
        # bpftrace prints objects as lines
        ret = {}
        for l in cmd_obj.stdout.split('\n'):
            if not l.strip():
                continue
            one = _json.loads(l)
            if one.get('type') != 'map':
                continue
            for k, v in one["data"].items():
                if k.startswith('@'):
                    k = k.lstrip('@')
                ret[k] = v
        return ret
    return cmd_obj


def rand_port(stype=socket.SOCK_STREAM):
    """
    Get a random unprivileged port.
    """
    return rand_ports(1, stype)[0]


def rand_ports(count, stype=socket.SOCK_STREAM):
    """
    Get a unique set of random unprivileged ports.
    """
    sockets = []
    ports = []

    try:
        for _ in range(count):
            s = socket.socket(socket.AF_INET6, stype)
            sockets.append(s)
            s.bind(("", 0))
            ports.append(s.getsockname()[1])
    finally:
        for s in sockets:
            s.close()

    return ports


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)


def wait_file(fname, test_fn, sleep=0.005, deadline=5, encoding='utf-8'):
    """
    Wait for file contents on the local system to satisfy a condition.
    test_fn() should take one argument (file contents) and return whether
    condition is met.
    """
    end = time.monotonic() + deadline

    with open(fname, "r", encoding=encoding) as fp:
        while True:
            if test_fn(fp.read()):
                break
            fp.seek(0)
            if time.monotonic() > end:
                raise TimeoutError("Wait for file contents failed", fname)
            time.sleep(sleep)