#! /usr/bin/env avocado-runner-avocado-classless

# SPDX-License-Identifier: GPL-2.0-or-later
#
# Copyright Red Hat
# Author: David Gibson <david@gibson.dropbear.id.au>

"""
Test A Simple Socket Transport

tasst/exesite.py - Manage simulated network sites for testing
"""


import contextlib
import ipaddress
import json

import avocado
from avocado.utils.process import CmdError
from avocado_classless.test import (
    assert_eq, assert_eq_unordered, assert_in, assert_raises, test_output
)

from tasst.typecheck import typecheck, typecheck_default


class SiteProcess(contextlib.AbstractContextManager):
    """
    A background process running on a Site
    """

    def __init__(self, site, cmd, subp, *,
                 ignore_status, context_timeout, pidfile):
        self.site = typecheck(site, Site)
        self.cmd = typecheck(cmd, str)
        self.subproc = typecheck(subp, avocado.utils.process.SubProcess)
        self.ignore_status = typecheck(ignore_status, bool)
        self.context_timeout = float(context_timeout)
        self.pidfile = typecheck_default(pidfile, str, None)
        self.pid = None

        if pidfile is not None:
            site.require_cmds('cat', 'kill')

    def __enter__(self):
        self.subproc.start()
        if self.pidfile is not None:
            # Wait for the PID file to be written
            pidstr = None
            while not pidstr:
                pidstr = self.site.output(f'cat {self.pidfile}',
                                          ignore_status=True)
            self.pid = int(pidstr)
        return self

    def __exit__(self, *exc_details):
        if self.pid is not None and self.subproc.poll() is None:
            self.site.fg(f'kill -TERM {self.pid}')
        result = self.subproc.run(timeout=self.context_timeout)
        if not self.ignore_status and result.exit_status != 0:
            siteinfo = f'[{self.site.name} site]'
            raise avocado.utils.process.CmdError(self.cmd, result, siteinfo)

    def run(self, **kwargs):
        return self.subproc.run(**kwargs)


class Site(contextlib.AbstractContextManager):
    """
    A (usually virtual or simulated) location where we can execute
    commands and configure networks.

    """

    def __init__(self, name):
        self.name = name   # For debugging

    def hostify(self, cmd, **kwargs):
        raise NotImplementedError

    def __enter__(self):
        return self

    def __exit__(self, *exc_details):
        pass

    def output(self, cmd, strip_trail_nl=False, **kwargs):
        kwargs['strip_trail_nl'] = strip_trail_nl
        cmd, kwargs = self.hostify(cmd, **kwargs)
        return avocado.utils.process.system_output(cmd, **kwargs)

    def fg(self, cmd, **kwargs):
        cmd, kwargs = self.hostify(cmd, **kwargs)
        return avocado.utils.process.system(cmd, **kwargs)

    def subprocess(self, cmd, **kwargs):
        cmd, kwargs = self.hostify(cmd, **kwargs)
        return avocado.utils.process.SubProcess(cmd, **kwargs)

    def bg(self, cmd, *,
           context_timeout=1.0, ignore_status=False, pidfile=None, **kwargs):
        subproc = self.subprocess(cmd, **kwargs)
        return SiteProcess(self, cmd, subproc,
                           context_timeout=context_timeout,
                           ignore_status=ignore_status, pidfile=pidfile)

    @contextlib.contextmanager
    def daemon(self, cmd, *, pidfile, **kwargs):
        self.require_cmds('cat', 'kill')
        self.fg(cmd, **kwargs)
        yield
        pid = int(self.output(f'cat {pidfile}'))
        self.fg(f'kill -TERM {pid}')

    def require_cmds(self, *cmds):
        missing = [c for c in cmds
                   if self.fg(f'type {c}', ignore_status=True) != 0]
        if missing:
            raise avocado.TestCancel(
                f"Missing commands {', '.join(missing)} on {self.name}")

    def ifs(self):
        self.require_cmds('ip')
        info = json.loads(self.output('ip -j link show'))
        return [i['ifname'] for i in info]

    def ifup(self, ifname, *addrs, dad=None):
        self.require_cmds('ip')
        if dad == 'disable':
            self.fg(f'sysctl net.ipv6.conf.{ifname}.accept_dad=0', sudo=True)
        elif dad == 'optimistic':
            self.fg(f'sysctl net.ipv6.conf.{ifname}.optimistic_dad=1',
                    sudo=True)
        elif dad is not None:
            raise ValueError

        for a in addrs:
            if not isinstance(a, ipaddress.IPv4Interface) \
               and not isinstance(a, ipaddress.IPv6Interface):
                raise TypeError
            self.fg(f'ip addr add {a.with_prefixlen} dev {ifname}', sudo=True)

        self.fg(f'ip link set {ifname} up', sudo=True)

    def addrinfos(self, ifname, **criteria):
        self.require_cmds('ip')
        info = json.loads(self.output(f'ip -j addr show {ifname}'))
        assert len(info) == 1  # We specified a specific interface

        ais = list(ai for ai in info[0]['addr_info'])
        for key, value in criteria.items():
            ais = [ai for ai in ais if key in ai and ai[key] == value]

        return ais

    def addrs(self, ifname, **criteria):
        self.require_cmds('ip')
        # Return just the parsed, non-tentative addresses
        return [ipaddress.ip_interface(f'{ai["local"]}/{ai["prefixlen"]}')
                for ai in self.addrinfos(ifname, **criteria)
                if 'tentative' not in ai]

    def mtu(self, ifname):
        self.require_cmds('ip')
        (info,) = json.loads(self.output(f'ip -j link show {ifname}'))
        return info['mtu']

    def addr_wait(self, ifname, **criteria):
        while True:
            addrs = self.addrs(ifname, **criteria)
            if addrs:
                return addrs

    def _routes(self, ipv, **criteria):
        routes = json.loads(self.output(f'ip -j -{ipv} route'))
        for key, value in criteria.items():
            routes = [r for r in routes if key in r and r[key] == value]

        return routes

    def routes4(self, **criteria):
        return self._routes('4', **criteria)

    def routes6(self, **criteria):
        return self._routes('6', **criteria)


def test_site(sitefn):
    def test_true(s):
        with s as site:
            site.fg('true')

    def test_false(s):
        with s as site:
            assert_raises(CmdError, site.fg, 'false')

    def test_echo(s):
        with s as site:
            msg = 'Hello tasst'
            out = site.output(f'echo {msg}')
            assert_eq(out, msg.encode('utf-8') + b'\n')

    def test_timeout(s):
        with s as site:
            site.fg('sleep infinity', timeout=0.1, ignore_status=True)

    def test_bg_true(s):
        with s as site:
            with site.bg('true'):
                pass

    def test_bg_false(s):
        with s as site:
            def run_false():
                with site.bg('false'):
                    pass
            assert_raises(CmdError, run_false)

    def test_bg_echo(s):
        msg = 'Hello tasst'
        with s as site:
            with site.bg(f'echo {msg}') as proc:
                res = proc.run()
        assert_eq(res.stdout, msg.encode('utf-8') + b'\n')

    def test_bg_timeout(s):
        with s as site:
            with site.bg('sleep infinity', ignore_status=True) as proc:
                proc.run(timeout=0.1)

    def test_bg_context_timeout(s):
        with s as site:
            with site.bg('sleep infinity', context_timeout=0.1,
                         ignore_status=True):
                pass

    def test_has_lo(s):
        with s as site:
            assert_in('lo', site.ifs())

    def test_lo_addrs(s):
        expected = [ipaddress.ip_interface(a)
                    for a in ['127.0.0.1/8', '::1/128']]
        with s as site:
            assert_eq_unordered(site.addrs('lo'), expected)

    def test_lo_mtu(s):
        with s as site:
            assert_eq(site.mtu('lo'), 65536)

    return test_output(test_true, test_false, test_echo, test_timeout,
                       test_bg_true, test_bg_false, test_bg_echo,
                       test_bg_timeout, test_bg_context_timeout,
                       test_has_lo, test_lo_addrs, test_lo_mtu)(sitefn)


def test_isolated_site(sitefn):
    def test_isolated_net(s):
        with s as site:
            assert_eq(site.ifs(), ['lo'])

    return test_output(test_isolated_net)(test_site(sitefn))


class RealHost(Site):
    """Represents the host on which the tests are running (as opposed
    to some simulated host created by the tests)

    """

    def __init__(self):
        super().__init__('REAL_HOST')

    def hostify(self, cmd, *, sudo=False, **kwargs):
        assert not sudo, "BUG: Shouldn't run commands with privilege on host"
        return cmd, kwargs


REAL_HOST = RealHost()


@test_site
def real_host():
    return REAL_HOST