Refactor Socket::Range

There was really no reason for it to inherit from Array. Also adds a few
more specs and gets coverage up to a more respectable percentage.

lib/rex/socket/range_walker.rb

@@ -13,12 +13,31 @@ module Socket
 # show-stoppingly inefficient when storing a bunch of non-consecutive
 # addresses, which should be a somewhat unusual case.
 #
+# @example
+#   r = RangeWalker.new("10.1,3.1-7.1-255")
+#   r.include?("10.3.7.255") #=> true
+#   r.length #=> 3570
+#   r.each do |addr|
+#     # do something with the address
+#   end
 ###
 class RangeWalker
 
+  # The total number of IPs within the range
   #
+  # @return [Fixnum]
+  attr_reader :length
+
+  # for backwards compatibility
+  alias :num_ips :length
+
+  # A list of the {Range ranges} held in this RangeWalker
+  # @return [Array]
+  attr_reader :ranges
+
   # Initializes a walker instance using the supplied range
   #
+  # @param parseme [RangeWalker,String]
   def initialize(parseme)
     if parseme.is_a? RangeWalker
       @ranges = parseme.ranges.dup
@@ -33,6 +52,7 @@ class RangeWalker
   #
   # This is basically only useful for determining if a range can be parsed
   #
+  # @return (see #parse)
   def self.parse(parseme)
     self.new.parse(parseme)
   end
@@ -41,19 +61,22 @@ class RangeWalker
   # Turn a human-readable range string into ranges we can step through one address at a time.
   #
   # Allow the following formats:
-  #	"a.b.c.d e.f.g.h"
+  #   "a.b.c.d e.f.g.h"
-  #	"a.b.c.d, e.f.g.h"
+  #   "a.b.c.d, e.f.g.h"
-  # where each chunk is CIDR notation, (e.g. '10.1.1.0/24') or a range in nmap format (see expand_nmap)
+  # where each chunk is CIDR notation, (e.g. '10.1.1.0/24') or a range in nmap format (see {#expand_nmap})
   #
   # OR this format
-  #	"a.b.c.d-e.f.g.h"
+  #   "a.b.c.d-e.f.g.h"
   # where a.b.c.d and e.f.g.h are single IPs and the second must be
   # bigger than the first.
   #
+  # @param parseme [String]
+  # @return [self]
+  # @return [false] if +parseme+ cannot be parsed
   def parse(parseme)
     return nil if not parseme
     ranges = []
-    parseme.split(', ').map{ |a| a.split(' ') }.flatten.each { |arg|
+    parseme.split(', ').map{ |a| a.split(' ') }.flatten.each do |arg|
       opts = {}
 
       # Handle IPv6 first (support ranges, but not CIDR)
@@ -67,7 +90,7 @@ class RangeWalker
 
           return false unless Rex::Socket.is_ipv6?(addr)
           addr = Rex::Socket.addr_atoi(addr)
-          ranges.push [addr, addr, true, opts]
+          ranges.push(Range.new(addr, addr, true, opts))
           next
         end
 
@@ -77,13 +100,14 @@ class RangeWalker
         addr2, scope_id = addrs[0].split('%')
         ( opts[:scope_id] ||= scope_id ) if scope_id
 
-        return false if not (Rex::Socket.is_ipv6?(addr1) and Rex::Socket.is_ipv6?(addr2))
+        # Both have to be IPv6 for this to work
+        return false unless (Rex::Socket.is_ipv6?(addr1) && Rex::Socket.is_ipv6?(addr2))
 
         # Handle IPv6 ranges in the form of 2001::1-2001::10
         addr1 = Rex::Socket.addr_atoi(addr1)
         addr2 = Rex::Socket.addr_atoi(addr2)
 
-        ranges.push [addr1, addr2, true, opts]
+        ranges.push(Range.new(addr1, addr2, true, opts))
         next
 
       # Handle IPv4 CIDR
@@ -116,19 +140,20 @@ class RangeWalker
         # Then it's a domain name and we should send it on to addr_atoi
         # unmolested to force a DNS lookup.
         begin
-          Rex::Socket.addr_atoi_list(arg).each { |addr| ranges.push [addr, addr, false, opts] }
+          ranges += Rex::Socket.addr_atoi_list(arg).map { |a| Range.new(a, a, false, opts) }
         rescue Resolv::ResolvError, ::SocketError, Errno::ENOENT
           return false
         end
 
       # Handle IPv4 ranges
       elsif arg =~ /^([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})-([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})$/
+
         # Then it's in the format of 1.2.3.4-5.6.7.8
         # Note, this will /not/ deal with DNS names, or the fancy/obscure 10...1-10...2
         begin
-          addrs = [Rex::Socket.addr_atoi($1), Rex::Socket.addr_atoi($2)]
+          start, stop = Rex::Socket.addr_atoi($1), Rex::Socket.addr_atoi($2)
-          return false if addrs[0] > addrs[1] # The end is greater than the beginning.
+          return false if start > stop # The end is greater than the beginning.
-          ranges.push [addrs[0], addrs[1], false, opts]
+          ranges.push(Range.new(start, stop, false, opts))
         rescue Resolv::ResolvError, ::SocketError, Errno::ENOENT
           return false
         end
@@ -139,7 +164,7 @@ class RangeWalker
           expanded.each { |r| ranges.push(r) }
         end
       end
-    }
+    end
 
     # Remove any duplicate ranges
     ranges = ranges.uniq
@@ -150,51 +175,58 @@ class RangeWalker
   #
   # Resets the subnet walker back to its original state.
   #
+  # @return [self]
   def reset
     return false if not valid?
     @curr_range = 0
-    @curr_addr = @ranges[0][0]
+    @curr_addr = @ranges.first.start
     @length = 0
-    @ranges.each { |r| @length += r[1] - r[0] + 1 }
+    @ranges.each { |r| @length += r.length }
+
+    self
   end
 
-  #
   # Returns the next IP address.
   #
+  # @return [String] The next address in the range
   def next_ip
     return false if not valid?
-    if (@curr_addr > @ranges[@curr_range][1])
+    if (@curr_addr > @ranges[@curr_range].stop)
       if (@curr_range >= @ranges.length - 1)
         return nil
       end
       @curr_range += 1
-      @curr_addr = @ranges[@curr_range][0]
+      @curr_addr = @ranges[@curr_range].start
     end
-    addr = Rex::Socket.addr_itoa(@curr_addr, @ranges[@curr_range][2])
+    addr = Rex::Socket.addr_itoa(@curr_addr, @ranges[@curr_range].ipv6)
 
-    if @ranges[@curr_range][3][:scope_id]
+    if @ranges[@curr_range].options[:scope_id]
-      addr = addr + '%' + @ranges[@curr_range][3][:scope_id]
+      addr = addr + '%' + @ranges[@curr_range].options[:scope_id]
     end
 
     @curr_addr += 1
     return addr
   end
 
+  alias :next :next_ip
+
+  # Whether this RangeWalker's ranges are valid
   def valid?
-    (@ranges and not @ranges.empty?)
+    (@ranges && !@ranges.empty?)
   end
 
-  #
   # Returns true if the argument is an ip address that falls within any of
   # the stored ranges.
   #
+  # @return [true] if this RangeWalker contains +addr+
+  # @return [false] if not
   def include?(addr)
     return false if not @ranges
     if (addr.is_a? String)
       addr = Rex::Socket.addr_atoi(addr)
     end
     @ranges.map { |r|
-      if r[0] <= addr and addr <= r[1]
+      if addr.between?(r.start, r.stop)
         return true
       end
     }
@@ -202,35 +234,45 @@ class RangeWalker
   end
 
   #
-  # Returns true if this RangeWalker includes all of the addresses in the
+  # Returns true if this RangeWalker includes *all* of the addresses in the
   # given RangeWalker
   #
-  def include_range?(range_walker)
+  # @param other [RangeWalker]
-    return false if ((not @ranges) or @ranges.empty?)
+  def include_range?(other)
-    return false if not range_walker.ranges
+    return false if (!@ranges || @ranges.empty?)
+    return false if !other.ranges || other.ranges.empty?
 
-    range_walker.ranges.all? do |start, stop|
+    # Check that all the ranges in +other+ fall within at least one of
-      ranges.any? do |self_start, self_stop|
+    # our ranges.
-        r = (self_start..self_stop)
+    other.ranges.all? do |other_range|
-        r.include?(start) and r.include?(stop)
+      ranges.any? do |range|
+        other_range.start.between?(range.start, range.stop) && other_range.stop.between?(range.start, range.stop)
       end
     end
   end
 
   #
   # Calls the given block with each address. This is basically a wrapper for
-  # #next_ip
+  # {#next_ip}
   #
+  # @return [self]
   def each(&block)
     while (ip = next_ip)
       block.call(ip)
     end
+    reset
+
+    self
   end
 
   #
-  # Returns an array with one element, a Range defined by the given CIDR
+  # Returns an Array with one element, a {Range} defined by the given CIDR
   # block.
   #
+  # @see Rex::Socket.cidr_crack
+  # @param arg [String] A CIDR range
+  # @return [Range]
+  # @return [false] if +arg+ is not valid CIDR notation
   def expand_cidr(arg)
     start,stop = Rex::Socket.cidr_crack(arg)
     if !start or !stop
@@ -373,27 +415,43 @@ class RangeWalker
     return ranges
   end
 
-  #
-  # The total number of IPs within the range
-  #
-  attr_reader :length
-
-  # for backwards compatibility
-  alias :num_ips :length
-
-  attr_reader :ranges
-
 end
 
-class Range < Array # :nodoc: all
+# A range of IP addresses
-  def start; self[0]; end
+class Range
-  def stop;  self[1]; end
+  def initialize(start=nil, stop=nil, ipv6=nil, options=nil)
-  def ipv6;  self[2]; end
+    @start = start
-  def options; self[3]; end
+    @stop = stop
-  def start=(val); self[0] = val; end
+    @ipv6 = ipv6
-  def stop=(val);  self[1] = val; end
+    @options = options
-  def ipv6=(val);  self[2] = val; end
+  end
-  def options=(val); self[3] = val; end
+
+  # Compare attributes with +other+
+  # @param other [Range]
+  def ==(other)
+    (other.start == start && other.stop == stop && other.ipv6 == ipv6 && other.options == options)
+  end
+
+  # The number of addresses in this Range
+  def length
+    stop - start + 1
+  end
+
+  #@!attribute start
+  #   The first address in this range, as a number
+  #   @return [Fixnum]
+  attr_accessor :start
+  #@!attribute stop
+  #   The last address in this range, as a number
+  #   @return [Fixnum]
+  attr_accessor :stop
+  #@!attribute ipv6
+  #   Whether this Range contains IPv6 or IPv4 addresses
+  #   @return [Boolean]
+  attr_accessor :ipv6
+  #@!attribute options
+  #   @return [Hash]
+  attr_accessor :options
 end
 
 end

spec/lib/rex/socket/range_walker_spec.rb

@@ -38,6 +38,11 @@ describe Rex::Socket::RangeWalker do
       it { should_not be_valid }
     end
 
+    context "with an IPv6 address range containing a scope" do
+      let(:args) { "fe80::1%lo-fe80::100%lo" }
+      it { should be_valid }
+    end
+
     it "should handle single ipv6 addresses" do
       walker = Rex::Socket::RangeWalker.new("::1")
       walker.should be_valid
@@ -51,6 +56,12 @@ describe Rex::Socket::RangeWalker do
       walker.next_ip.should == "10.1.1.1"
     end
 
+    context "with mulitple ranges" do
+      let(:args) { "1.1.1.1-2 2.1-2.2.2 3.1-2.1-2.1 " }
+      it { should be_valid }
+      it { should have(8).addresses }
+      it { should include("1.1.1.1") }
+    end
 
     it "should handle ranges" do
       walker = Rex::Socket::RangeWalker.new("10.1.1.1-2")
@@ -141,14 +152,27 @@ describe Rex::Socket::RangeWalker do
         walker.length.should == (2**(32-bits))
       end
     end
+  end
+
+  describe '#each' do
+    let(:args) { "10.1.1.1-2,2,3 10.2.2.2" }
 
     it "should yield all ips" do
-      walker = Rex::Socket::RangeWalker.new("10.1.1.1,2,3")
       got = []
       walker.each { |ip|
         got.push ip
       }
-      got.should == ["10.1.1.1", "10.1.1.2", "10.1.1.3"]
+      got.should == ["10.1.1.1", "10.1.1.2", "10.1.1.3", "10.2.2.2"]
+    end
+
+  end
+
+  describe '#include_range?' do
+    let(:args) { "10.1.1.*" }
+
+    it "returns true for a sub-range" do
+      other = described_class.new("10.1.1.1-255")
+      walker.should be_include_range(other)
     end
 
   end