subnet walker action

git-svn-id: file:///home/svn/incoming/trunk@2977 4d416f70-5f16-0410-b530-b9f4589650da

lib/rex/socket.rb

@@ -25,6 +25,7 @@ module Socket
 	require 'rex/socket/comm/local'
 	
 	require 'rex/socket/switch_board'
+	require 'rex/socket/subnet_walker'
 
 	##
 	#
@@ -90,10 +91,20 @@ module Socket
 		return to_sockaddr(host, 0)[4,4]
 	end
 
+	#
+	# Resolves a host to a network-byte order ruby integer.
+	#
 	def self.resolv_nbo_i(host)
 		return resolv_nbo(host).unpack('N')[0]
 	end
 
+	#
+	# Resolves a host to a dotted address.
+	#
+	def self.resolv_to_dotted(host)
+		Resolv.getaddress(host)
+	end
+
 	#
 	# Converts a netmask (255.255.255.240) into a bitmask (28).  This is the
 	# lame kid way of doing it.

lib/rex/socket/subnet_walker.rb

@@ -0,0 +1,63 @@
+require 'rex/socket'
+
+module Rex
+module Socket
+
+###
+#
+# SubnetWalker
+# ------------
+#
+# This class provides an interface to enumerating a subnet with a supplied
+# netmask.
+#
+###
+class SubnetWalker
+
+	def initialize(subnet, netmask)
+		self.subnet  = Socket.resolv_to_dotted(subnet)
+		self.netmask = Socket.resolv_to_dotted(netmask)
+
+		reset
+	end
+
+	#
+	# Resets the subnet walker back to its original state.
+	#
+	def reset
+		self.curr_ip     = self.subnet.split('.')
+		self.num_ips     = (1 << (32 - Socket.net2bitmask(self.netmask).to_i)) - 1
+		self.curr_ip_idx = 0
+	end
+
+	#
+	# Returns the next IP address.
+	#
+	def next_ip
+		if (curr_ip_idx > num_ips)
+			return nil
+		end
+
+		if (curr_ip_idx > 0)
+			self.curr_ip[3] = (curr_ip[3].to_i + 1) % 256
+			self.curr_ip[2] = (curr_ip[2].to_i + 1) % 256 if (curr_ip[3] == 0)
+			self.curr_ip[1] = (curr_ip[1].to_i + 1) % 256 if (curr_ip[2] == 0)
+			self.curr_ip[0] = (curr_ip[0].to_i + 1) % 256 if (curr_ip[1] == 0)
+		end
+
+		self.curr_ip_idx += 1
+	
+		self.curr_ip.join('.')
+	end
+
+	attr_reader :subnet, :netmask, :num_ips
+
+protected
+
+	attr_writer   :subnet, :netmask, :num_ips
+	attr_accessor :curr_ip, :curr_ip_idx
+
+end
+
+end
+end

lib/rex/socket/subnet_walker.rb.ut.rb

@@ -0,0 +1,28 @@
+#!/usr/bin/ruby
+
+$:.unshift(File.join(File.dirname(__FILE__), '..', '..'))
+
+require 'test/unit'
+require 'rex/socket/subnet_walker'
+
+class Rex::Socket::SubnetWalker::UnitTest < Test::Unit::TestCase
+
+	Klass = Rex::Socket::SubnetWalker
+
+	def test_walker
+		s = Klass.new('10.0.0.0', '255.255.255.0')
+
+		0.upto(255) { |x|
+			assert_equal('10.0.0.' + x.to_s, s.next_ip)
+		}
+		assert_nil(s.next_ip)
+
+		s.reset
+
+		0.upto(255) { |x|
+			assert_equal('10.0.0.' + x.to_s, s.next_ip)
+		}
+		assert_nil(s.next_ip)
+	end
+
+end