#! /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/address.py - Address allocation helpers
#
# Copyright Red Hat
# Author: David Gibson <david@gibson.dropbear.id.au>

import ipaddress

import avocado

from tasst import Tasst


# Loopback addresses, for convenience
LOOPBACK4 = ipaddress.ip_address('127.0.0.1')
LOOPBACK6 = ipaddress.ip_address('::1')

# Documentation test networks defined by RFC 5737
TEST_NET_1 = ipaddress.ip_network('192.0.2.0/24')
TEST_NET_2 = ipaddress.ip_network('198.51.100.0/24')
TEST_NET_3 = ipaddress.ip_network('203.0.113.0/24')

# Documentation test network defined by RFC 3849
TEST_NET6 = ipaddress.ip_network('2001:db8::/32')
# Some subnets of that for our usage
TEST_NET6_TASST_A = ipaddress.ip_network('2001:db8:9a55:aaaa::/64')
TEST_NET6_TASST_B = ipaddress.ip_network('2001:db8:9a55:bbbb::/64')
TEST_NET6_TASST_C = ipaddress.ip_network('2001:db8:9a55:cccc::/64')


class IpiAllocator:
    DEFAULT_NETS = [TEST_NET_1, TEST_NET6_TASST_A]

    def __init__(self, *nets):
        if not nets:
            nets = self.DEFAULT_NETS

        self.nets = [ipaddress.ip_network(n) for n in nets]
        self.hostses = [n.hosts() for n in self.nets]

    def next_ipis(self):
        addrs = [next(h) for h in self.hostses]
        return [ipaddress.ip_interface('{}/{}'.format(a, n.prefixlen))
                     for a, n in zip(addrs, self.nets)]


class IpiAllocatorTasst(Tasst):
    """
    :avocado: tags=meta
    """

    COUNT = 12  # Allocate this many addresses
    NETS = None # From these networks (None means use default)

    def test(self):
        if self.NETS is not None:
            ipa = IpiAllocator(*self.NETS)
            nets = self.NETS
        else:
            ipa = IpiAllocator()
            nets = IpiAllocator.DEFAULT_NETS

        addrsets = [set() for i in range(len(nets))]
        for i in range(self.COUNT):
            addrs = ipa.next_ipis()
            # Check we got as many addresses as expected
            self.assertEquals(len(addrs), len(nets))
            for s, a, n in zip(addrsets, addrs, nets):
                # Check the addresses belong to the right network
                self.assertEquals(a.network, ipaddress.ip_network(n))
                s.add(a)

        print(addrsets)
        # Check the addresses are unique
        for s in addrsets:
            self.assertEquals(len(s), self.COUNT)

class IpiAllocatorCustomTasst(IpiAllocatorTasst):
    """
    :avocado: tags=meta
    """

    NETS = ['10.55.0.0/16', '192.168.55.0/24', 'fd00:9a57:a000::/48']