Merge pull request #84 from garo/feature/allocate

Add allocate() method

4 files changed, 110 insertions, 0 deletions

diff --git a/lib/ipaddress/ipv4.rb b/lib/ipaddress/ipv4.rb
index ade1851..70d1a51 100644
--- a/lib/ipaddress/ipv4.rb
+++ b/lib/ipaddress/ipv4.rb
@@ -90,6 +90,7 @@ module IPAddress;
       # 32 bits interger containing the address
       @u32 = (@octets[0]<< 24) + (@octets[1]<< 16) + (@octets[2]<< 8) + (@octets[3])
       
+      @allocator = 0
     end # def initialize
 
     #
@@ -1080,6 +1081,36 @@ module IPAddress;
     end
 
     #
+    # Allocates a new ip from the current subnet. Optional skip parameter
+    # can be used to skip addresses.
+    #
+    # Will raise StopIteration exception when all addresses have been allocated
+    #
+    # Example:
+    #
+    #    ip = IPAddress("10.0.0.0/24")
+    #    ip.allocate
+    #      #=> "10.0.0.1/24"
+    #    ip.allocate
+    #      #=> "10.0.0.2/24"
+    #    ip.allocate(2)
+    #      #=> "10.0.0.5/24"
+    #
+    #
+    # Uses an internal @allocator which tracks the state of allocated
+    # addresses.
+    #
+    def allocate(skip=0)
+        @allocator += 1 + skip
+
+        next_ip = network_u32+@allocator
+        if next_ip > broadcast_u32+1
+            raise StopIteration
+        end
+        self.class.parse_u32(network_u32+@allocator, @prefix)
+    end
+
+    #
     # private methods
     #
     private
diff --git a/lib/ipaddress/ipv6.rb b/lib/ipaddress/ipv6.rb
index adc6243..f427954 100644
--- a/lib/ipaddress/ipv6.rb
+++ b/lib/ipaddress/ipv6.rb
@@ -103,6 +103,7 @@ module IPAddress;
       end
 
       @prefix = Prefix128.new(netmask ? netmask : 128)
+      @allocator = 0
 
     end # def initialize
 
@@ -664,6 +665,36 @@ module IPAddress;
     def self.parse_hex(hex, prefix=128)
       self.parse_u128(hex.hex, prefix)
     end
+
+    #
+    # Allocates a new ip from the current subnet. Optional skip parameter
+    # can be used to skip addresses.
+    #
+    # Will raise StopIteration exception when all addresses have been allocated
+    #
+    # Example:
+    #
+    #    ip = IPAddress("10.0.0.0/24")
+    #    ip.allocate
+    #      #=> "10.0.0.1/24"
+    #    ip.allocate
+    #      #=> "10.0.0.2/24"
+    #    ip.allocate(2)
+    #      #=> "10.0.0.5/24"
+    #
+    #
+    # Uses an internal @allocator which tracks the state of allocated
+    # addresses.
+    #
+    def allocate(skip=0)
+        @allocator += 1 + skip
+
+        next_ip = network_u128+@allocator
+        if next_ip > broadcast_u128
+            raise StopIteration
+        end
+        self.class.parse_u128(next_ip, @prefix)
+    end
     
     private
 
diff --git a/test/ipaddress/ipv4_test.rb b/test/ipaddress/ipv4_test.rb
index f39d686..19e0450 100644
--- a/test/ipaddress/ipv4_test.rb
+++ b/test/ipaddress/ipv4_test.rb
@@ -627,6 +627,30 @@ class IPv4Test < Minitest::Test
     assert_equal "192.168.200.0/24", ip.to_string
   end
 
+  def test_allocate_addresses
+    ip = @klass.new("10.0.0.0/24")
+    ip1 = ip.allocate
+    ip2 = ip.allocate
+    ip3 = ip.allocate
+    assert_equal "10.0.0.1/24", ip1.to_string
+    assert_equal "10.0.0.2/24", ip2.to_string
+    assert_equal "10.0.0.3/24", ip3.to_string
+  end
+
+  def test_allocate_can_skip_addresses
+    ip = @klass.new("10.0.0.0/24")
+    ip1 = ip.allocate(2)
+    assert_equal "10.0.0.3/24", ip1.to_string
+  end
+
+  def test_allocate_will_raise_stopiteration
+    ip = @klass.new("10.0.0.0/30")
+    ip.allocate(3)
+    assert_raises (StopIteration) do
+      ip.allocate
+    end
+  end
+
 end # class IPv4Test
 
   
diff --git a/test/ipaddress/ipv6_test.rb b/test/ipaddress/ipv6_test.rb
index 120aafb..9d7736d 100644
--- a/test/ipaddress/ipv6_test.rb
+++ b/test/ipaddress/ipv6_test.rb
@@ -282,6 +282,30 @@ class IPv6Test < Minitest::Test
     assert_equal expected, arr
   end
 
+  def test_allocate_addresses
+    ip = @klass.new("2001:db8::4/125")
+    ip1 = ip.allocate
+    ip2 = ip.allocate
+    ip3 = ip.allocate
+    assert_equal "2001:db8::1", ip1.compressed
+    assert_equal "2001:db8::2", ip2.compressed
+    assert_equal "2001:db8::3", ip3.compressed
+  end
+
+  def test_allocate_can_skip_addresses
+    ip = @klass.new("2001:db8::4/125")
+    ip1 = ip.allocate(2)
+    assert_equal "2001:db8::3", ip1.compressed
+  end
+
+  def test_allocate_will_raise_stopiteration
+    ip = @klass.new("2001:db8::4/125")
+    ip.allocate(6)
+    assert_raises (StopIteration) do
+      ip.allocate
+    end
+  end
+
   def test_method_compare
     ip1 = @klass.new("2001:db8:1::1/64")
     ip2 = @klass.new("2001:db8:2::1/64")