diff options
Diffstat (limited to 'vicn/resource/ip/prefix_tree.py')
| -rw-r--r-- | vicn/resource/ip/prefix_tree.py | 156 |
1 files changed, 15 insertions, 141 deletions
diff --git a/vicn/resource/ip/prefix_tree.py b/vicn/resource/ip/prefix_tree.py index 34af1d14..1fec2d47 100644 --- a/vicn/resource/ip/prefix_tree.py +++ b/vicn/resource/ip/prefix_tree.py @@ -16,134 +16,7 @@ # limitations under the License. # -from socket import inet_pton, inet_ntop, AF_INET6 -from struct import unpack, pack -from abc import ABCMeta - -class PrefixTreeException(Exception): pass -class NotEnoughAddresses(PrefixTreeException): pass -class UnassignablePrefix(PrefixTreeException): pass - -class Prefix(metaclass=ABCMeta): - - def __init__(self, ip_address, prefix_size=None): - if not prefix_size: - ip_address, prefix_size = ip_address.split('/') - prefix_size = int(prefix_size) - if isinstance(ip_address, str): - ip_address = self.aton(ip_address) - self.ip_address = ip_address - self.prefix_size = prefix_size - - def __contains__(self, obj): - #it can be an IP as a integer - if isinstance(obj, int): - obj = type(self)(obj, self.MAX_PREFIX_SIZE) - #Or it's an IP string - if isinstance(obj, str): - #It's a prefix as 'IP/prefix' - if '/' in obj: - split_obj = obj.split('/') - obj = type(self)(split_obj[0], int(split_obj[1])) - else: - obj = type(self)(obj, self.MAX_PREFIX_SIZE) - - return self._contains_prefix(obj) - - @classmethod - def mask(cls): - mask_len = cls.MAX_PREFIX_SIZE//8 #Converts from bits to bytes - mask = 0 - for step in range(0,mask_len): - mask = (mask << 8) | 0xff - return mask - - def _contains_prefix(self, prefix): - assert isinstance(prefix, type(self)) - return (prefix.prefix_size >= self.prefix_size and - prefix.ip_address >= self.first_prefix_address() and - prefix.ip_address <= self.last_prefix_address()) - - #Returns the first address of a prefix - def first_prefix_address(self): - return self.ip_address & (self.mask() << (self.MAX_PREFIX_SIZE-self.prefix_size)) - - def last_prefix_address(self): - return self.ip_address | (self.mask() >> self.prefix_size) - - def limits(self): - return self.first_prefix_address(), self.last_prefix_address() - - def __str__(self): - return "{}/{}".format(self.ntoa(self.first_prefix_address()), self.prefix_size) - - def __eq__(self, other): - return (self.first_prefix_address() == other.first_prefix_address() and - self.prefix_size == other.prefix_size) - - def __hash__(self): - return hash(str(self)) - - def __iter__(self): - return self.get_iterator() - - #Iterates by steps of prefix_size, e.g., on all available /31 in a /24 - def get_iterator(self, prefix_size=None): - if prefix_size is None: - prefix_size=self.MAX_PREFIX_SIZE - assert (prefix_size >= self.prefix_size and prefix_size<=self.MAX_PREFIX_SIZE) - step = 2**(self.MAX_PREFIX_SIZE - prefix_size) - for ip in range(self.first_prefix_address(), self.last_prefix_address()+1, step): - yield type(self)(ip, prefix_size) - -class Inet4Prefix(Prefix): - - MAX_PREFIX_SIZE = 32 - - @classmethod - def aton(cls, address): - ret = 0 - components = address.split('.') - for comp in components: - ret = (ret << 8) + int(comp) - return ret - - @classmethod - def ntoa(cls, address): - components = [] - for _ in range(0,4): - components.insert(0,'{}'.format(address % 256)) - address = address >> 8 - return '.'.join(components) - -class Inet6Prefix(Prefix): - - MAX_PREFIX_SIZE = 128 - - @classmethod - def aton (cls, address): - prefix, suffix = unpack(">QQ", inet_pton(AF_INET6, address)) - return (prefix << 64) | suffix - - @classmethod - def ntoa (cls, address): - return inet_ntop(AF_INET6, pack(">QQ", address >> 64, address & ((1 << 64) -1))) - - #skip_internet_address: skip a:b::0, as v6 often use default /64 prefixes - def get_iterator(self, prefix_size=None, skip_internet_address=None): - if skip_internet_address is None: - #We skip the internet address if we iterate over Addresses - if prefix_size is None: - skip_internet_address = True - #But not if we iterate over prefixes - else: - skip_internet_address = False - it = super().get_iterator(prefix_size) - if skip_internet_address: - next(it) - return it - -###### PREFIX TREE ###### +from netmodel.model.type import Inet4Prefix, Inet6Prefix class PrefixTree: @@ -160,21 +33,21 @@ class PrefixTree: self.full = False - def find_prefix(self, prefix_size): + def find_prefix(self, prefix_len): ret, lret, rret = [None]*3 - if prefix_size > self.prefix.prefix_size and not self.full: + if prefix_len > self.prefix.prefix_len and not self.full: if self.left is None: - lret = self.prefix_cls(self.prefix.first_prefix_address(), self.prefix.prefix_size+1) + lret = self.prefix_cls(self.prefix.first_prefix_address(), self.prefix.prefix_len+1) else: - lret = self.left.find_prefix(prefix_size) + lret = self.left.find_prefix(prefix_len) if self.right is None: - rret = self.prefix_cls(self.prefix.last_prefix_address(), self.prefix.prefix_size+1) + rret = self.prefix_cls(self.prefix.last_prefix_address(), self.prefix.prefix_len+1) else: - rret = self.right.find_prefix(prefix_size) + rret = self.right.find_prefix(prefix_len) #Now we look for the longer prefix to assign - if not lret or (rret and rret.prefix_size > lret.prefix_size): + if not lret or (rret and rret.prefix_len > lret.prefix_len): ret = rret else: ret = lret @@ -183,9 +56,9 @@ class PrefixTree: def assign_prefix(self, prefix): assert prefix in self.prefix - if prefix.prefix_size > self.prefix.prefix_size: + if prefix.prefix_len > self.prefix.prefix_len: #Existing prefix on the left - lp = self.prefix_cls(self.prefix.first_prefix_address(), self.prefix.prefix_size+1) + lp = self.prefix_cls(self.prefix.first_prefix_address(), self.prefix.prefix_len+1) #It's on the left branch if prefix in lp: if not self.left: @@ -193,7 +66,7 @@ class PrefixTree: self.left.assign_prefix(prefix) #It's on the right branch else: - rp = self.prefix_cls(self.prefix.last_prefix_address(), self.prefix.prefix_size+1) + rp = self.prefix_cls(self.prefix.last_prefix_address(), self.prefix.prefix_len+1) if not self.right: self.right = PrefixTree(rp) self.right.assign_prefix(prefix) @@ -202,13 +75,14 @@ class PrefixTree: else: raise RuntimeError("And you may ask yourself, well, how did I get here?") - def get_prefix(self, prefix_size): - ret = self.find_prefix(prefix_size) + def get_prefix(self, prefix_len): + ret = self.find_prefix(prefix_len) if not ret: raise NotEnoughAddresses #find_prefix returns the size of the largest available prefix in our space #not necessarily the one we asked for - ret.prefix_size = prefix_size + ret.ip_address = ret.first_prefix_address() + ret.prefix_len = prefix_len self.assign_prefix(ret) return ret |