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

from avocado_classless.test import assert_eq, test_output

from tasst.address import IpiAllocator, TEST_NET_1
from tasst.nstool import unshare_site


DHCLIENT = '/sbin/dhclient'


@contextlib.contextmanager
def dhclient(site, ifname, ipv='4'):
    site.require_cmds(DHCLIENT)

    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 {pidfile} -lf {leasefile} {ifname}'
        with site.daemon(f'{DHCLIENT} {opts}', sudo=True, pidfile=pidfile):
            yield


def test_dhcp(ifname, expected_addr, gateway, mtu):
    def dec(setupfn):
        def test_addr(setup):
            with setup as site, dhclient(site, ifname):
                (actual_addr,) = site.addrs(ifname, family='inet',
                                            scope='global')
                assert_eq(actual_addr.ip, expected_addr)

        def test_route(setup):
            with setup as site, dhclient(site, ifname):
                (defroute,) = site.routes4(dst='default')
                assert_eq(ipaddress.ip_address(defroute['gateway']), gateway)

        def test_mtu(setup):
            with setup as site, dhclient(site, ifname):
                assert_eq(site.mtu(ifname), mtu)

        return test_output(test_addr, test_route, test_mtu)(setupfn)

    return dec


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


@contextlib.contextmanager
def setup_dhcpd_common(ifname, server_ifname):
    with unshare_site('client', '-Un') as client, \
         unshare_site('server', '-n', parent=client, sudo=True) as server:
        server.require_cmds(DHCPD)

        client.veth(ifname, server_ifname, server)

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


@test_dhcp(IFNAME, CLIENT_IP4.ip, SERVER_IP4.ip, 1500)
@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'-f -d -4 -cf {confpath} -lf {leasepath} -pf {pidfile}'
        server.fg(f'{DHCPD} -t {opts}')  # test config
        with server.bg(f'{DHCPD} {opts}', sudo=True, ignore_status=True,
                       pidfile=pidfile):
            # Configure the client
            client.ifup('lo')
            yield client