#! /usr/bin/python3

# SPDX-License-Identifier: GPL-2.0-or-later
#
# tasst - Test A Simple Socket Transport
#         library of test helpers for passt & pasta
#
# tasst/site.py - Manage simulated network sites for testing
#
# Copyright Red Hat
# Author: David Gibson <david@gibson.dropbear.id.au>

import contextlib
import ipaddress
import json

import avocado
from avocado.utils.process import CmdError

from tasst import Tasst


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 __enter__(self):
        raise NotImplementedError

    def __exit__(self, *exc_details):
        raise NotImplementedError

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

    def fg(self, cmd, **kwargs):
        self.output(cmd, **kwargs)

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

    def require_cmds(self, *cmds):
        missing = [c for c in cmds
                   if self.fg('type {}'.format(c), ignore_status=True) != 0]
        if missing:
            raise avocado.TestCancel("Missing commands {} on {}"
                                     .format(', '.join(missing), 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('sysctl net.ipv6.conf.{}.accept_dad=0'.format(ifname), sudo=True)
        elif dad == 'optimistic':
            self.fg('sysctl net.ipv6.conf.{}.optimistic_dad=1'.format(ifname), 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('ip addr add {} dev {}'.format(a.with_prefixlen, ifname), sudo=True)

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

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

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

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

        return ais

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

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


class SiteTasst(Tasst):
    """
    Basic tests for executing commands on sites

    :avocado: disable
    :avocado: tags=meta
    """
    timeout = 1.0

    # Derived classes must redefine this
    def setup_site(self):
        raise NotImplementedError("{} must implement setup_site() method".format(type(self).__name__))

    def test_true(self):
        with self.setup_site() as site:
            site.fg('true')

    def test_false(self):
        with self.setup_site() as site:
            self.assertRaises(CmdError, site.fg, 'false')

    def test_echo(self):
        with self.setup_site() as site:
            s = 'Hello tasst'
            out = site.output('echo {}'.format(s))
            self.assertEquals(out, s.encode('utf-8'))

    def test_bg_true(self):
        with self.setup_site() as site:
            with site.bg('true') as proc:
                status = proc.wait()
                self.assertEquals(status, 0)

    def test_bg_false(self):
        with self.setup_site() as site:
            with site.bg('false') as proc:
                status = proc.wait()
                self.assertNotEquals(status, 0)

    def test_has_lo(self):
        with self.setup_site() as site:
            self.assertIn('lo', site.ifs())

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

    def test_lo_mtu(self):
        with self.setup_site() as site:
            self.assertEqual(site.mtu('lo'), 65536)


class IsolatedSiteTasst(SiteTasst):
    """
    Test a site with isolated network (loopback only)

    :avocado: disable
    :avocado: tags=meta
    """

    def test_isolated_net(self):
        with self.setup_site() as site:
            self.assertEquals(site.ifs(), ['lo'])


# Represents the host on which the tests are running, as opposed to
# some simulated host created by the tests
class RealHost(Site):
    def __init__(self):
        super().__init__('REAL_HOST')

    def __enter__(self):
        return self

    def __exit__(self, *exc_details):
        pass

    def output(self, cmd, sudo=False, **kwargs):
        assert not sudo, "BUG: Shouldn't run commands with privilege on host"
        return avocado.utils.process.system_output(cmd, **kwargs)

    def fg(self, cmd, sudo=False, **kwargs):
        assert not sudo, "BUG: Shouldn't run commands with privilege on host"
        return avocado.utils.process.system(cmd, **kwargs)

    @contextlib.contextmanager
    def bg(self, cmd, sudo=False, **kwargs):
        assert not sudo, "BUG: Shouldn't run commands with privilege on host"
        try:
            proc = avocado.utils.process.SubProcess(cmd, **kwargs)
            proc.start()
            yield proc
        finally:
            if proc.poll() is None:
                proc.stop()


REAL_HOST = RealHost()


class RealHostTasst(SiteTasst):
    def setup_site(self):
        return REAL_HOST