Expose ContainsNet Add GetNetworks (#13)

* Expose ContainsNet Add GetNetworks

* Replace containsNet with ContainsNet

* fix wording

Author: smikulcik

ipset.go

@@ -114,6 +114,15 @@ func (s *IPSet) GetIPs(limit int) (ips []net.IP) {
 	return
 }
 
+// GetNetworks retrieves a list of all networks included in the ipTree
+func (s *IPSet) GetNetworks() []*net.IPNet {
+	networks := []*net.IPNet{}
+	s.tree.walk(func(node *ipTree) {
+		networks = append(networks, node.net)
+	})
+	return networks
+}
+
 // Intersection computes the set intersect between this IPSet and another one
 // It returns a new set which is the intersection.
 func (s *IPSet) Intersection(set1 *IPSet) (interSect *IPSet) {

ipset_test.go

@@ -1,6 +1,7 @@
 package netaddr
 
 import (
+	"fmt"
 	"math/big"
 	"math/rand"
 	"net"
@@ -418,3 +419,15 @@ func TestIPSetAllocateDeallocate(t *testing.T) {
 	assert.Equal(t, 51036, len(ips))
 	assert.Equal(t, []error{}, available.tree.validate())
 }
+
+func TestGetNetworks(t *testing.T) {
+	s := &IPSet{}
+	assert.Equal(t, []*net.IPNet{}, s.GetNetworks())
+	s.InsertNet(Ten24)
+	assert.Equal(t, "[10.0.0.0/24]", fmt.Sprintf("%s", s.GetNetworks()))
+	ten25, _ := ParseNet("10.0.0.0/25")
+	s.RemoveNet(ten25)
+	assert.Equal(t, "[10.0.0.128/25]", fmt.Sprintf("%s", s.GetNetworks()))
+	s.Remove(ParseIP("10.0.0.129"))
+	assert.Equal(t, "[10.0.0.128/32 10.0.0.130/31 10.0.0.132/30 10.0.0.136/29 10.0.0.144/28 10.0.0.160/27 10.0.0.192/26]", fmt.Sprintf("%s", s.GetNetworks()))
+}

iptree.go

@@ -42,7 +42,7 @@ func (t *ipTree) trimLeft(top *ipTree) *ipTree {
 		return nil
 	}
 
-	if containsNet(top.net, t.net) {
+	if ContainsNet(top.net, t.net) {
 		return t.left.trimLeft(top)
 	}
 	t.setRight(t.right.trimLeft(top))
@@ -55,7 +55,7 @@ func (t *ipTree) trimRight(top *ipTree) *ipTree {
 		return nil
 	}
 
-	if containsNet(top.net, t.net) {
+	if ContainsNet(top.net, t.net) {
 		return t.right.trimRight(top)
 	}
 	t.setLeft(t.left.trimRight(top))
@@ -71,11 +71,11 @@ func (t *ipTree) insert(newNode *ipTree) *ipTree {
 		return newNode
 	}
 
-	if containsNet(t.net, newNode.net) {
+	if ContainsNet(t.net, newNode.net) {
 		return t
 	}
 
-	if containsNet(newNode.net, t.net) {
+	if ContainsNet(newNode.net, t.net) {
 		// Replace the current top node and trim the tree
 		newNode.setLeft(t.left.trimLeft(newNode))
 		newNode.setRight(t.right.trimRight(newNode))
@@ -98,10 +98,10 @@ func (t *ipTree) contains(newNode *ipTree) bool {
 		return false
 	}
 
-	if containsNet(t.net, newNode.net) {
+	if ContainsNet(t.net, newNode.net) {
 		return true
 	}
-	if containsNet(newNode.net, t.net) {
+	if ContainsNet(newNode.net, t.net) {
 		return false
 	}
 	if bytes.Compare(newNode.net.IP, t.net.IP) < 0 {
@@ -161,10 +161,10 @@ func (t *ipTree) removeNet(net *net.IPNet) (top *ipTree) {
 	}
 
 	top = t
-	if containsNet(net, t.net) {
+	if ContainsNet(net, t.net) {
 		// Remove the current node
 		top = t.remove()
-	} else if containsNet(t.net, net) {
+	} else if ContainsNet(t.net, net) {
 		diff = netDifference(t.net, net)
 		t.net = diff[0]
 		for _, n := range diff[1:] {

net_utils.go

@@ -92,9 +92,9 @@ func BroadcastAddr(n *net.IPNet) net.IP {
 	return broadcast
 }
 
-// containsNet returns true if net2 is a subset of net1. To be clear, it
+// ContainsNet returns true if net2 is a subset of net1. To be clear, it
 // returns true if net1 == net2 also.
-func containsNet(net1, net2 *net.IPNet) bool {
+func ContainsNet(net1, net2 *net.IPNet) bool {
 	// If the two networks are different IP versions, return false
 	if len(net1.IP) != len(net2.IP) {
 		return false
@@ -117,11 +117,11 @@ func netDifference(a, b *net.IPNet) (result []*net.IPNet) {
 	}
 
 	// If b contains a then the difference is empty
-	if containsNet(b, a) {
+	if ContainsNet(b, a) {
 		return
 	}
 	// If a doesn't contain b then the difference is equal to a
-	if !containsNet(a, b) {
+	if !ContainsNet(a, b) {
 		return []*net.IPNet{a}
 	}