# David.Clark (a) CIBER-research.eu 2012-12-20 # CIBER Research Ltd. # version 10 hybrid for both clarity and optimisation, + bug fix. # this code is online at # http://code.activestate.com/recipes/576479/ """ Create a digital trie of ipranges(CIDR); for a given ip determine membership of smallest CIDR (if any) by digital search. 'Instead of basing a search method on comparisons between keys, we can make use of their representation as a sequence of digits or alphabetic characters.' Knuth. Art of Computer Programming vol 3, 2e 1998 (6.3 p492) """ def bytes2bits(): """Create a table of bit values for int(0) to int(256).""" # derived from Python Cookbook 2e 2005. bytes = [None] * 256 for n in xrange(256): bits = [] q = n # q is modified by bitshift for b in xrange(0,8): if (q & 128): bits.append(1) else: bits.append(0) q <<= 1 bytes[n] = tuple(bits) # make immutable return tuple(bytes) # One-time call to build table; rebind the funtion name to its result. bytes2bits = bytes2bits() class ProgramError(ValueError): """This is a bug""" class CIDRnode(object): def __init__(self, cidr=None): self.cidr = cidr self.left = None self.right = None def get_net(self): if self.cidr: return self.cidr.split('/')[0] else: return None network = property(get_net, doc='IPv4 dotted quad') def get_bits(self): if self.cidr: return int(self.cidr.split('/')[1]) else: return None netbits = property(get_bits, doc='IPv4 netmask bits') class CIDRtrie(object): """ take a shortcut by using a list for the first octet as any CIDR will be no greater than class A. Thereafter work down a binary tree of ip network addresses. If a node has a CIDR record stack the result, but continue down the tree while there are links. (Not all Nodes have CIDR records.). At leif node pop (the best) result. """ # we need some definitions before we can __init__() @staticmethod def byte2bit(ip): for q in ip: q = int(q) for b in xrange(0,8): yield 1 if (q & 128) else 0 q <<= 1 def add_cidr(self, cidr): """Build the trie. For clarity this is the non-optimised version.""" c = cidr.split('/') network = c[0].split('.') nm_bit = int(c[1]) classA = int(network[0]) subtree = self.root[classA] if not subtree: subtree = CIDRnode() self.root[classA] = subtree nm_bit -= 7 # leave a bit for test at top of loop network_bits = self.byte2bit(network[1:]) for nextbit in network_bits: if nm_bit == 1: overwrite = subtree.cidr subtree.cidr = cidr return overwrite # expect to return None nm_bit -= 1 if nextbit is 0: if not subtree.left: subtree.left = CIDRnode() subtree = subtree.left continue elif nextbit is 1: if not subtree.right: subtree.right = CIDRnode() subtree = subtree.right continue else: raise ProgramError( 'Tried to bud %s, bitten by bug, fell out of Tree'% cidr) subtree.cidr = cidr def __init__(self, ipranges=None): self.root = [None] * 256 # A forest of (x.0.0.0/8) 'class A' addresses if ipranges: for cidr in ipranges: self.add_cidr(cidr) def get_cidr(self, ip): """This is very similar to add_cidr but inline code improves performance by around 10% and the cost of building the lookup table is amortised over serveral million lookups gaining another 10%.""" ip = ip.split('.') subtree = self.root[int(ip[0])] # subtree = CIDRnode if subtree is None: return None results = [None] for quad in ip[1:]: quad_bits = bytes2bits[int(quad)] for nextbit in quad_bits: if subtree.cidr: results.append(subtree.cidr) if subtree.left and nextbit is 0: subtree = subtree.left continue elif subtree.right and nextbit is 1: subtree = subtree.right continue else: return results.pop() return subtree.cidr #### Command Line Processing #### def test(): index = {} networks = ( ('', 'network 1'), ('' , 'network 2'), ('' , 'network 3'), ('','network 4'), ('' , 'single_ip'),) ipaddr = ('', '', '', '', '', '', '', '', '') trie = CIDRtrie() for cidr, name in networks: index[cidr] = name overwrite = trie.add_cidr(cidr) if overwrite: print('WARNING overwriting %s with %s'% (overwrite, cidr)) for ip in ipaddr: cidr = trie.get_cidr(ip) if cidr: name = index[cidr] print('%s is in network %s (%s)'% (ip, name, cidr)) else: print('network not known for %s'% ip) if __name__ == '__main__': #only when run from cmd line test()