#! /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/dhcp.py - Helpers for testing DHCP
#
# Copyright Red Hat
# Author: David Gibson <david@gibson.dropbear.id.au>

import contextlib
import ipaddress
import os

from tasst import Tasst, TasstSubData
from tasst.address import IpiAllocator, TEST_NET_1
from tasst.nstool import UnshareSite
from tasst.site import Site
from tasst.typing import typecheck


class DhcpTasstInfo:
    def __init__(self, site, ifname, addr, gw, mtu):
        self.site = typecheck(site, Site)
        self.ifname = typecheck(ifname, str)
        self.addr = typecheck(addr, ipaddress.IPv4Address)
        self.gw = typecheck(gw, ipaddress.IPv4Address)
        self.mtu = typecheck(mtu, int)

        site.require_cmds('ip')


class BaseDhcpTasst(Tasst):
    """
    Test DHCP behaviour.

    :avocado: disable
    """

    DHCLIENT = '/sbin/dhclient'

    def setup_dhcp(self):
        raise NotImplementedError("{} must implement setup_dhcp() method".format(type(self).__name__))

    @contextlib.contextmanager
    def dhclient(self):
        with self.setup_dhcp() as dti:
            dti = typecheck(dti, DhcpTasstInfo)

            dti.site.require_cmds(self.DHCLIENT, 'cat', 'kill')

            pidfile = os.path.join(self.workdir, 'dhclient.pid')
            leasefile = os.path.join(self.workdir, 'dhclient.leases')

            # We need '-nc' because we may be running with
            # capabilities but not UID 0.  Without -nc dhclient drops
            # capabilities before invoking dhclient-script, so it's
            # unable to actually configure the interface
            dti.site.fg('{} -4 -v -nc -pf {} -lf {} {}'
                        .format(self.DHCLIENT, pidfile, leasefile, dti.ifname), sudo=True)

            try:
                yield dti
            finally:
                pid = int(dti.site.output('cat {}'.format(pidfile)))
                dti.site.fg('kill {}'.format(pid))

    def test_addr(self):
        with self.dhclient() as dti:
            (addr,) = dti.site.addrs(dti.ifname, family='inet', scope='global')
            self.assertEquals(addr.ip, dti.addr)

    def test_route(self):
        with self.dhclient() as dti:
            (defroute,) = dti.site.routes4(dst='default')
            self.assertEquals(ipaddress.ip_address(defroute['gateway']), dti.gw)

    def test_mtu(self):
        with self.dhclient() as dti:
            self.assertEquals(dti.site.mtu(dti.ifname), dti.mtu)


class MetaDhcpTasst(BaseDhcpTasst):
    """Ugly workaround for
    https://github.com/avocado-framework/avocado/issues/5680.
    Explicitly apply the "meta" tag to inherited tests

    :avocado: disable
    :avocado: tags=meta

    """

    def test_addr(self):
        super().test_addr()

    def test_route(self):
        super().test_route()

    def test_mtu(self):
        super().test_mtu()


class DhcpdTasst(MetaDhcpTasst):
    DHCPD = 'dhcpd'
    SUBNET = TEST_NET_1

    @contextlib.contextmanager
    def setup_dhcp(self):
        ifname = 'clientif'
        server_ifname = 'serverif'

        with UnshareSite(type(self).__name__ + '.client', '-Un') as client, \
             UnshareSite(type(self).__name__ + '.server', '-n',
                         parent=client, sudo=True) as server:

            server.require_cmds(self.DHCPD)

            client.veth(ifname, server_ifname, server)

            # Configure the DHCP server
            ipa = IpiAllocator(self.SUBNET)
            (server_ip4,) = ipa.next_ipis()
            (client_ip4,) = ipa.next_ipis()

            confpath = os.path.join(self.workdir, 'dhcpd.conf')
            open(confpath, 'w').write('''
            subnet {} netmask {} {{
                option routers {};
                range {} {};
            }}
            '''.format(self.SUBNET.network_address, self.SUBNET.netmask,
                       server_ip4.ip, client_ip4.ip, client_ip4.ip))
            pidpath = os.path.join(self.workdir, 'dhcpd.pid')
            leasepath = os.path.join(self.workdir, 'dhcpd.leases')
            open(leasepath, 'w').write('')

            server.ifup('lo')
            server.ifup(server_ifname, server_ip4)

            opts = ('-f -d -4 -cf {} -lf {} -pf {}'.format(confpath, leasepath, pidpath))
            server.fg('{} -t {}'.format(self.DHCPD, opts)) # test config
            with server.bg('{} {}'.format(self.DHCPD, opts), sudo=True) as dhcpd:
                # Configure the client
                client.ifup('lo')

                yield DhcpTasstInfo(client, ifname, client_ip4.ip, server_ip4.ip, 1500)

                pid = int(open(pidpath).read())
                server.fg('kill {}'.format(pid))
                status = dhcpd.wait()