Welcome, guest | Sign In | My Account | Store | Cart
# slais-sysweb (a) ucl.ac.uk	2008-09-02

# UCL Centre for Publishing
# version 8 hybrid for both clarity and optimisation.

""" 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):
            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]
            return None
    network = property(get_net, doc='IPv4 dotted quad')

    def get_bits(self):
        if self.cidr:
            return int(self.cidr.split('/')[1])
            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__()
    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
            elif nextbit is 1:
                if not subtree.right:
                    subtree.right = CIDRnode()
                subtree = subtree.right
                raise ProgramError(
                    'Tried to bud %s, bitten by bug, fell out of Tree'% 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:

    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:
                if subtree.left and nextbit is 0:
                    subtree = subtree.left
                elif subtree.right and nextbit is 1:
                    subtree = subtree.right
                    return results.pop()

####    Command Line Processing    ####
def test():
    index = {}
    networks = (
        ('', 'network 1'),
        ('' , 'network 2'),
        ('' , 'network 3'),
        ('','network 4'))
    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))
            print('network not known for %s'% ip)

if __name__ == '__main__':     #only when run from cmd line