ipadmin 0.2.1 → 0.2.2

data/lib/tree.rb

@@ -0,0 +1,735 @@
+=begin rdoc
+ Copyright (c) 2006 Dustin Spinhirne -  
+ Licensed under the same terms as Ruby, No Warranty is provided.
+=end
+
+
+module IPAdmin
+class Tree
+
+# instance variables
+# @all_f
+# @max_bits
+# @root
+# @version
+
+#==============================================================================#
+# attr_reader / attr_writer
+#==============================================================================#
+
+    # IP version for this tree (4 or 6)
+    attr_reader :version
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# initialize()
+#==============================================================================#
+
+# - Arguments:
+#   * Hash with the following fields:
+#       - :Version -- IP version - Integer
+#
+# Example:
+#   table = IPAdmin::CIDRTable.new(:Version => 4)
+#
+    def initialize(options)
+        if (!options.kind_of? Hash)
+            raise ArgumentError, "Expected Hash, but #{options.class} provided."
+        end
+
+        if ( options.has_key?(:Version) )
+            @version =  options[:Version]
+
+            if ( @version != 4 && @version != 6 )
+                raise "IP version should be either 4 or 6."
+            end
+        else
+            raise ArgumentError, "Missing argument: Version."
+        end
+
+        if (@version == 4)
+            @max_bits = 32
+            @all_f = 2**@max_bits - 1
+        else
+            @max_bits = 128
+            @all_f = 2**@max_bits - 1
+        end
+
+        # root of our ordered IP tree
+        @root = []
+
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# add()
+#==============================================================================#
+
+# Add an IPAdmin::CIDR object to the tree.
+#
+# - Arguments:
+#   * CIDR object
+#
+# - Returns:
+#   * nothing
+#
+# Example:
+#   tree.add(cidr)
+#
+    def add(object)
+
+        # validate object
+        if ( !object.kind_of?(IPAdmin::CIDR) )
+            raise ArgumentError, "IPAdmin::CIDR object " +
+                                 "required but #{object.class} provided."
+        end
+
+        if ( object.version != @version )
+            raise "IP version #{object.version} is incompatible with " +
+                  "Tree IP version #{@version}."
+        end
+
+        net_struct = IPAdmin.create_net_struct(object)
+
+        # search tree for it's home branch
+        search_results = find_home(net_struct, @root)
+        home_struct = search_results[0] if (search_results)
+        home_branch = home_struct.subnets if (search_results)
+        home_branch = @root if (!home_branch)
+
+        # make sure we dont have a duplicate
+        if (home_struct)
+            if ( (net_struct.network == home_struct.network) && 
+                 (net_struct.netmask == home_struct.netmask) )
+                raise "Attempted to add #{object.desc()} multiple times."
+            end
+        end
+
+        # now that we have our branch location, check that other entries
+        # of this branch are not subnets of our new object. if they are
+        # then make them a branch of our new object
+        home_branch.each do |entry|
+            is_contained = contains(net_struct, entry)
+
+            if (is_contained)
+                net_struct.subnets.push(entry)
+            end
+        end
+
+        net_struct.subnets.each do |entry|
+            home_branch.delete(entry)
+        end
+
+        # add new object as an ordered entry to home_branch
+        add_to_branch(net_struct,home_branch)
+
+        return(nil)
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# collapse()
+#==============================================================================#
+
+# Collapse (supernet) all subnets of the tree,
+# and return a new tree. The supernetting of subnets will only occur in
+# situations where the newly created supernet will not result in the
+# 'creation' of additional space. For example the following blocks
+# (192.168.0.0/24, 192.168.1.0/24, and 192.168.2.0/24) would merge into
+# 192.168.0.0/23 and 192.168.2.0/24 rather than into 192.168.0.0/22.
+#
+# - Arguments:
+#   * none
+#
+# - Returns:
+#   * Tree object
+#
+# Example:
+#   new_tree = tree.collapse()
+#
+    def collapse()
+            tree = IPAdmin::Tree.new(:Version => @version)
+            branch = IPAdmin.merge(:List => @root, :NetStruct => 1)
+            dumped = dump_branch(branch)
+            dumped.each do |entry|
+                net_struct = entry[:NetStruct]
+                network = IPAdmin.unpack_ip_addr(:Integer => net_struct.network,
+                                                 :Version => @version)
+                netmask = IPAdmin.unpack_ip_netmask(:Integer => net_struct.netmask,
+                                                    :Version => @version)
+                cidr = IPAdmin::CIDR.new(:CIDR => "#{network}/#{netmask}")
+                tree.add(cidr)
+            end
+        return(tree)
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# dump
+#==============================================================================#
+
+# Dump the contents of this tree.
+#
+# - Arguments:
+#   * none
+#
+# - Returns:
+#   * ordered array of hashes with the following fields: 
+#       - :Object => CIDR object
+#       - :Depth => (depth level in tree)
+# Example:
+#   dumped = tree.dump()
+#
+    def dump()
+        list = []
+        dumped = dump_branch(@root)
+
+        dumped.each do |entry|
+            depth = entry[:Depth]
+            net_struct = entry[:NetStruct]
+            object = net_struct.object
+            list.push({:Depth => depth, :Object => object})
+        end
+
+        return(list)
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# exists?()
+#==============================================================================#
+
+# Has an IPAdmin::CIDR object already been added to the tree?
+#
+# - Arguments:
+#   * CIDR object
+#
+# - Returns:
+#   * true or false
+#
+# Example:
+#   added = tree.exists?(cidr)
+#
+    def exists?(object)
+
+        # validate object
+        if ( !object.kind_of?(IPAdmin::CIDR) )
+            raise ArgumentError, "IPAdmin::CIDR object " +
+                                 "required but #{object.class} provided."
+        end
+
+        if ( !object.version == @version )
+            raise "IP version #{object.version} is incompatible with " +
+                  "Tree IP version #{@version}."
+        end
+
+        net_struct = IPAdmin.create_net_struct(object)
+        home_struct,home_branch = find_home(net_struct, @root)
+
+        # check for match
+        found = false
+        if (home_struct)  
+            if ( (net_struct.network == home_struct.network) && 
+                 (net_struct.netmask == home_struct.netmask) )
+                found = true
+            end
+        end
+
+
+        return(found)
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# find()
+#==============================================================================#
+
+# Find the longest matching branch of our tree to which an 
+# IPAdmin::CIDR object belongs. Useful for performing 'routing table' lookups.
+#
+# - Arguments:
+#   * CIDR object
+#
+# - Returns:
+#   * CIDR object, or nil
+#
+# Example:
+#   longest_match = tree.find(ip)
+#
+    def find(object)
+        if ( !object.kind_of?(IPAdmin::CIDR) )
+            raise ArgumentError, "IPAdmin::CIDR object " +
+                                 "required but #{object.class} provided."
+        end
+
+        if ( object.version != @version )
+            raise "IP version #{object.version} is incompatible with " +
+                  "Tree IP version #{@version}."
+        end
+
+        net_struct = IPAdmin.create_net_struct(object)
+        home_struct,home_branch = find_home(net_struct, @root)
+        found_obj = home_struct.object if (home_struct)
+
+        return(found_obj)
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# find_space()
+#==============================================================================#
+
+# Find Tree entries that can hold a subnet of size X. Returns only the smallest
+# matching subnets of each supernet. If neither IPCount or Subnet or provided,
+# then method returns all entries that can fit a single IP.
+#
+# - Arguments:
+#   * Hash with the following fields:
+#       - :IPCount -- Minimum number of IP's that should fit within subnet - Integer (optional)
+#       - :Limit -- Maximum entries to return - Integer (optional)
+#       - :Subnet -- Netmask (in bits) of space to find - Integer (optional)
+#
+# - Returns:
+#   * Array of CIDR objects
+#
+# - Notes:
+#    * :Subnet always takes precedence over :IPCount.
+#
+# Example:
+#   list = tree.find_space(:Subnet => 27)
+#   list = tree.find_space(:IPCount => 33, :Limit => 2)
+#
+    def find_space(options)
+        limit = nil
+        list = []
+
+        # validate options
+        if (!options.kind_of? Hash)
+            raise ArgumentError, "Expected Hash, but #{options.class} provided."
+        end
+
+        if ( options.has_key?(:Limit) )
+            limit = options[:Limit]
+        end
+        
+        if ( options.has_key?(:IPCount) )
+            subnet_size = IPAdmin.minimum_size(:IPCount =>options[:IPCount],
+                                               :Version => @version)
+        end
+        
+        if ( options.has_key?(:Subnet) )
+            subnet_size = options[:Subnet]
+        end
+
+        # set subnet_size if not already
+        if (!subnet_size)
+            if (@version == 4)
+                subnet_size = 32
+            else
+                subnet_size = 128
+            end
+        end
+        
+        # check that subnet_size is a valid size
+        if ( (subnet_size < 1) || (subnet_size > @max_bits) )
+                raise "#{subnet_size} is out of range for an " +
+                      "IP version #{@version} Tree."
+        end
+
+        # search
+        search_list = dump_branch(@root)
+        search_list.each do |entry|
+            depth = entry[:Depth]
+            net_struct = entry[:NetStruct]
+
+            # we only want leaf nodes of type IPAdmin::CIDR
+            if ( (net_struct.subnets.length == 0) &&
+                 (net_struct.object.kind_of?(IPAdmin::CIDR)) )
+
+                 if (subnet_size >= net_struct.object.bits)
+                    list.push(net_struct.object)
+                 end
+            end
+
+            if (limit && (list.length >= limit) )
+                break
+            end
+
+        end
+
+        return(list)
+
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# prune()
+#==============================================================================#
+
+# Remove an IPAdmin::CIDR object, and all child
+# objects from the tree.
+#
+# - Arguments:
+#   * CIDR object
+#
+# - Returns:
+#   * true or false
+#
+# Example:
+#   did_prune = tree.prune(ip)
+#
+    def prune(object)
+        if ( !object.kind_of?(IPAdmin::CIDR) )
+            raise ArgumentError, "IPAdmin::CIDR object " +
+                                 "required but #{object.class} provided."
+        end
+
+        if (object.version != @version )
+            raise "IP version #{object.version} is incompatible with " +
+                  "Tree IP version #{@version}."
+        end
+
+        net_struct = IPAdmin.create_net_struct(object)
+        home_struct,home_branch = find_home(net_struct, @root)
+
+        if(!home_struct)
+            raise "#{object.desc} could not be found."
+        end
+        
+        # remove if home_struct.object = object
+        pruned = false
+        if (home_struct.object = object)
+            home_branch.delete(home_struct)
+            pruned = true
+        end
+
+        return(pruned)
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# remove()
+#==============================================================================#
+
+# Remove a single CIDR object from the tree.
+#
+# - Arguments:
+#   * CIDR object
+#
+# - Returns:
+#   * true or false
+#
+# Example:
+#   did_remove = tree.remove(ip)
+#
+    def remove(object)
+        if ( !object.kind_of?(IPAdmin::CIDR) )
+            raise ArgumentError, "IPAdmin::CIDR object " +
+                                 "required but #{object.class} provided."
+        end
+
+        if (object.version != @version )
+            raise "IP version #{object.version} is incompatible with " +
+                  "Tree IP version #{@version}."
+        end
+
+        net_struct = IPAdmin.create_net_struct(object)
+        home_struct,home_branch = find_home(net_struct, @root)
+
+        # remove if home_struct.object = object
+        removed = false
+        if (home_struct.object = object)
+
+            # if we have children subnets, move them up one level
+            if (home_struct.subnets.length > 0)
+                home_struct.subnets.each do |entry|
+                    index = 0
+                    home_branch.each do
+                        if(entry.network < (home_branch[index]).network)
+                            break
+                        end
+                        index += 1
+                    end
+                    home_branch.insert(index, entry)
+                end
+            end
+
+            home_branch.delete(home_struct)
+            removed = true
+        end
+
+        return(removed)
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# show()
+#==============================================================================#
+
+# Print the tree in a nicely formatted string.
+#
+# - Arguments:
+#   * none
+#
+# - Returns:
+#   * String representing the contents of the tree
+#
+# Example:
+#   puts tree.show()
+#
+    def show()
+        printed = ""
+        list = dump_branch(@root)
+
+        list.each do |entry|
+            net_struct = entry[:NetStruct]
+            depth = entry[:Depth]
+            network = IPAdmin.unpack_ip_addr(:Integer => net_struct.network,
+                                             :Version => @version)
+            netmask = IPAdmin.unpack_ip_netmask(:Integer => net_struct.netmask,
+                                                :Version => @version)
+            network = IPAdmin.shorten(network) if (@version == 6)
+
+            if (depth == 0)
+                indent = ""
+            else
+                indent = " " * (depth*3)
+            end
+
+            printed << "#{indent}#{network}/#{netmask}\n"
+        end
+
+        return(printed)
+    end
+
+#======================================#
+#
+#======================================#
+
+
+
+
+# PRIVATE METHODS
+private
+
+#==============================================================================#
+# add_to_branch() private
+#==============================================================================#
+
+# Add Netstruct object to an array of NetStruct's
+#
+# - Arguments:
+#   * NetStruct, and array of NetStruct's in which to add it
+#
+# - Returns:
+#   * none
+#
+# Example:
+#   add_to_branch(net_struct,branch)
+#
+    def add_to_branch(net_struct,branch)
+        # make sure we have appropriate types 
+        if (!net_struct.kind_of?(IPAdmin::NetStruct) || 
+            !branch.kind_of?(Array) )
+            raise ArgumentError, "Incorrect argument types " +
+                                 "#{net_struct.class} and #{branch}."
+        end
+        
+        index = 0
+        branch.each do
+            if(net_struct.network < (branch[index]).network)
+                break
+            end
+            index += 1
+        end
+
+        branch.insert(index, net_struct)
+
+        return()
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# contains() private
+#==============================================================================#
+
+# Is the first NetStruct object the supernet of the second?
+#
+# - Arguments:
+#   * Two NetStruct objects
+#
+# - Returns:
+#   * true if struct1 contains or is equal to struct2, or false if not.
+#
+# Example:
+#   is_contained = contains(struct1,struct2)
+#
+    def contains(struct1,struct2)
+
+        # make sure we have appropriate types 
+        if (!struct1.kind_of?(IPAdmin::NetStruct) || 
+            !struct2.kind_of?(IPAdmin::NetStruct) )
+            raise ArgumentError, "Incorrect argument types " +
+                                 "#{struct1.class} and #{struct2.class}."
+        end
+
+        network1 = struct1.network
+        netmask1 = struct1.netmask
+        network2 = struct2.network
+        netmask2 = struct2.netmask
+        hostmask1 = netmask1 ^ @all_f
+
+        # if network1 == network2, then we have to go by netmask length
+        # else we can tell by or'ing network1 and network2 by hostmask1
+        # and comparing the results
+        is_contained = false
+        if (network1 == network2)
+            if (netmask2 >= netmask1)
+                is_contained = true
+            end
+        else
+            if ( (network1 | hostmask1) == (network2 | hostmask1) )
+                is_contained = true
+            end
+        end
+
+        return(is_contained)
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# dump_branch() private
+#==============================================================================#
+
+#  Dump contents of an Array of NetStruct objects
+#
+# - Arguments:
+#   * array of NetStruct objects, and an optional depth within the tree
+#
+# - Returns:
+#   * array of hashes in form: :NetStruct => IPAdmin::NetStruct object
+#                              :Depth => (depth level in tree)
+# Example: 
+#   dumped = dump_branch(branch)
+#
+    def dump_branch(branch,depth=nil)
+        list = []
+        depth = 0 if (!depth)
+
+        branch.each do |entry|
+            subnets = entry.subnets
+            list.push({:NetStruct => entry, :Depth => depth})
+
+            if (subnets.length > 0)
+                list.concat( dump_branch(subnets, (depth+1) ) )
+            end
+
+        end
+
+        return(list)
+    end
+
+#======================================#
+#
+#======================================#
+
+
+#==============================================================================#
+# find_home() private
+#==============================================================================#
+
+# Find the branch to which a NetStruct object belongs.
+#
+# - Arguments:
+#   * IPAdmin::NetStruct to find, and Array of
+#     IPAdmin::NetStruct's to search
+#
+# - Returns:
+#   * array of: container IPAdmin::NetStruct object,
+#               and container branch, or nil
+#
+# Example:
+#   branch = find_home(net_struct,branch)
+#
+    def find_home(net_struct,branch)
+        ret_val = nil
+
+        branch.each do |entry|
+            is_contained = contains(entry,net_struct)
+
+            if (is_contained)
+                home_struct = entry
+                home_branch = branch
+
+                if (home_struct.subnets.length > 0)
+                    search_results = find_home(net_struct,home_struct.subnets)
+                    if (search_results)
+                        home_struct = search_results[0] 
+                        home_branch = search_results[1]
+                    end
+                end
+
+                ret_val = [home_struct,home_branch]
+                break
+            end
+        end
+
+        return(ret_val)
+    end
+
+#======================================#
+#
+#======================================#
+
+end
+
+end # module IPAdmin
+__END__