Created
November 2, 2021 15:59
-
-
Save twisteroidambassador/9fa22eebb75fcd4835aa9344fcd57770 to your computer and use it in GitHub Desktop.
Why is picking a random IP subnet in Python so awkward?
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import ipaddress | |
import typing | |
import abc | |
IPNetworkType = typing.TypeVar('IPNetworkType', ipaddress.IPv4Network, ipaddress.IPv6Network) | |
class IPSubnetSequence(typing.Sequence[IPNetworkType]): | |
"""An indexable sequence of subnets of a given network. | |
This is like ipaddress.IPv{4,6}Network.subnets(), except that method returns a generator which can only be | |
consumed in order. This class works more like range(), or a list where entries are generated on demand, | |
so you can directly retrieve entries in the middle. Useful for picking random subnets, for example. | |
""" | |
def __init__( | |
self, | |
base_network: IPNetworkType, | |
/, | |
prefixlen_diff: typing.Optional[int] = None, | |
new_prefixlen: typing.Optional[int] = None, | |
): | |
"""Constructor. | |
One and only one of prefixlen_diff or new_prefixlen must be set. | |
Args: | |
base_network: network to be divided into subnets | |
prefixlen_diff: number of bits the prefix should be increased by | |
new_prefixlen: desired new prefix length of subnets | |
""" | |
self._base_network = base_network | |
for network_type in (ipaddress.IPv4Network, ipaddress.IPv6Network): | |
if isinstance(base_network, network_type): | |
self._network_type = network_type | |
break | |
else: | |
raise TypeError(f'{base_network!r} is not an IP network type') | |
if new_prefixlen is not None: | |
if prefixlen_diff is not None: | |
raise ValueError('Cannot set prefixlen_diff and new_prefixlen at the same time') | |
else: | |
if prefixlen_diff is None: | |
raise ValueError('Must set either prefixlen_diff or new_prefixlen') | |
if prefixlen_diff <= 0: | |
raise ValueError('prefixlen_diff must be > 0') | |
new_prefixlen = self._base_network.prefixlen + prefixlen_diff | |
if not self._base_network.prefixlen < new_prefixlen <= self._base_network.max_prefixlen: | |
raise ValueError(f'Prefix length {new_prefixlen} invalid for network type {self._network_type}') | |
self._new_prefixlen = new_prefixlen | |
prefixlen_diff = new_prefixlen - self._base_network.prefixlen | |
self._num_subnets = 2 ** prefixlen_diff | |
self._subnets_address_range = range( | |
int(self._base_network[0]), | |
int(self._base_network[-1]) + 1, | |
(int(self._base_network.hostmask) + 1) >> prefixlen_diff, | |
) | |
@property | |
def base(self): | |
"""The base network.""" | |
return self._base_network | |
@property | |
def new_prefixlen(self): | |
"""The prefix length of subnets.""" | |
return self._new_prefixlen | |
def __len__(self): | |
return self._num_subnets | |
@typing.overload | |
@abc.abstractmethod | |
def __getitem__(self, i: int) -> IPNetworkType: | |
... | |
@typing.overload | |
@abc.abstractmethod | |
def __getitem__(self, s: slice) -> typing.Sequence[IPNetworkType]: | |
... | |
def __getitem__(self, i): | |
if isinstance(i, int): | |
try: | |
subnet_address = self._subnets_address_range[i] | |
except IndexError: | |
raise IndexError('Index out of range') | |
return self._network_type((subnet_address, self._new_prefixlen)) | |
if isinstance(i, slice): | |
return [self.__getitem__(idx) for idx in range(*i.indices(self.__len__()))] | |
raise TypeError | |
if __name__ == '__main__': | |
# Example: picking 10 random IPv6 ULA networks | |
import random | |
ula_space = ipaddress.IPv6Network('fd00::/8') | |
ula_networks_seq = IPSubnetSequence(ula_space, new_prefixlen=48) | |
for _ in range(10): | |
print(random.choice(ula_networks_seq)) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment