Skip to content

Instantly share code, notes, and snippets.

@twisteroidambassador
Created November 2, 2021 15:59
Show Gist options
  • Save twisteroidambassador/9fa22eebb75fcd4835aa9344fcd57770 to your computer and use it in GitHub Desktop.
Save twisteroidambassador/9fa22eebb75fcd4835aa9344fcd57770 to your computer and use it in GitHub Desktop.
Why is picking a random IP subnet in Python so awkward?
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