#! /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

dhcp.py - Helpers for testing DHCP
"""

import contextlib
import ipaddress
import os
import tempfile

import exeter

from . import address, nstool


DHCLIENT = '/sbin/dhclient'


@contextlib.contextmanager
def dhclient(snh, ifname, ipv='4'):
    with tempfile.TemporaryDirectory() as tmpdir:
        pidfile = os.path.join(tmpdir, 'dhclient.pid')
        leasefile = os.path.join(tmpdir, '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
        opts = [f'-{ipv}', '-v', '-nc', '-pf', f'{pidfile}',
                '-lf', f'{leasefile}', f'{ifname}']
        snh.fg(f'{DHCLIENT}', *opts, capable=True)
        yield
        snh.fg(f'{DHCLIENT}', '-x', '-pf', f'{pidfile}', capable=True)


class DhcpTestScenario:
    def __init__(self, *, client, ifname, addr, gateway, mtu):
        self.client = client
        self.ifname = ifname
        self.addr = addr
        self.gateway = gateway
        self.mtu = mtu


def test_dhcp_addr(setup):
    with setup as scn, dhclient(scn.client, scn.ifname):
        (actual_addr,) = scn.client.addrs(scn.ifname,
                                          family='inet', scope='global')
        exeter.assert_eq(actual_addr.ip, scn.addr)


def test_dhcp_route(setup):
    with setup as scn, dhclient(scn.client, scn.ifname):
        (defroute,) = scn.client.routes4(dst='default')
        exeter.assert_eq(ipaddress.ip_address(defroute['gateway']),
                         scn.gateway)


def test_dhcp_mtu(setup):
    with setup as scn, dhclient(scn.client, scn.ifname):
        exeter.assert_eq(scn.client.mtu(scn.ifname), scn.mtu)


DHCP_TESTS = [test_dhcp_addr, test_dhcp_route, test_dhcp_mtu]


def dhcp_tests(setup):
    for t in DHCP_TESTS:
        testid = f'{setup.__qualname__}|{t.__qualname__}'
        exeter.register_pipe(testid, setup, t)


DHCPD = 'dhcpd'
SUBNET = address.TEST_NET_1
ipa = address.IpiAllocator(SUBNET)
(SERVER_IP4,) = ipa.next_ipis()
(CLIENT_IP4,) = ipa.next_ipis()
IFNAME = 'clientif'


@contextlib.contextmanager
def setup_dhcpd_common(ifname, server_ifname):
    with nstool.unshare_snh('client', '-Un') as client, \
         nstool.unshare_snh('server', '-n',
                            parent=client, capable=True) as server:
        client.veth(ifname, server_ifname, server)

        with tempfile.TemporaryDirectory() as tmpdir:
            yield (client, server, tmpdir)


@contextlib.contextmanager
def setup_dhcpd():
    server_ifname = 'serverif'

    with setup_dhcpd_common(IFNAME, server_ifname) as (client, server, tmpdir):
        # Configure dhcpd
        confpath = os.path.join(tmpdir, 'dhcpd.conf')
        open(confpath, 'w', encoding='UTF-8').write(
            f'''subnet {SUBNET.network_address} netmask {SUBNET.netmask} {{
            option routers {SERVER_IP4.ip};
            range {CLIENT_IP4.ip} {CLIENT_IP4.ip};
            }}'''
        )
        pidfile = os.path.join(tmpdir, 'dhcpd.pid')
        leasepath = os.path.join(tmpdir, 'dhcpd.leases')
        open(leasepath, 'wb').write(b'')

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

        opts = ['-f', '-d', '-4', '-cf', f'{confpath}',
                '-lf', f'{leasepath}', '-pf', f'{pidfile}']
        server.fg(f'{DHCPD}', '-t', *opts)  # test config
        with server.bg(f'{DHCPD}', *opts, capable=True, check=False) as dhcpd:
            # Configure the client
            client.ifup('lo')
            yield DhcpTestScenario(client=client, ifname=IFNAME,
                                   addr=CLIENT_IP4.ip,
                                   gateway=SERVER_IP4.ip, mtu=1500)
            dhcpd.terminate()


dhcp_tests(setup_dhcpd)