ip_addr.rb in ronin-support-0.3.0

- old
+ new

@@ -15,69 +15,75 @@
 #
 # You should have received a copy of the GNU Lesser General Public License
 # along with Ronin Support.  If not, see <http://www.gnu.org/licenses/>.
 #
 
+require 'ronin/extensions/resolv'
+require 'ronin/extensions/regexp'
+
 require 'ipaddr'
-require 'resolv'
 require 'strscan'
 require 'combinatorics/list_comprehension'
 
 class IPAddr
 
   include Enumerable
 
-  # A regular expression for matching IPv4 Addresses.
-  IPV4_REGEXP = /[\d]{1,3}\.[\d]{1,3}\.[\d]{1,3}\.[\d]{1,3}/
+  # Socket families and IP address masks
+  MASKS = {
+    Socket::AF_INET  => IN4MASK,
+    Socket::AF_INET6 => IN6MASK
+  }
 
-  # A regular expression for matching IPv6 Addresses.
-  IPV6_REGEXP = /:(:[0-9a-f]{1,4}){1,7}|([0-9a-f]{1,4}::?){1,7}[0-9a-f]{1,4}(:#{IPV4_REGEXP})?/
-
-  # A regular expression for matching IP Addresses.
-  REGEXP = /#{IPV4_REGEXP}|#{IPV6_REGEXP}/
-
   #
   # Extracts IP Addresses from text.
   #
   # @param [String] text
   #   The text to scan for IP Addresses.
   #
-  # @param [Symbol] version
-  #   The version of IP Address to scan for (`:ipv4` or `:ipv6`).
+  # @param [Integer, Symbol] version
+  #   The version of IP Address to scan for (`4`, `6`, `:v4` or `:v6`).
   #
   # @yield [ip]
   #   The given block will be passed each extracted IP Address.
   #
   # @yieldparam [String] ip
   #   An IP Address from the text.
   #
   # @return [Array<String>]
   #   The IP Addresses found in the text.
   #
+  # @example
+  #   IPAddr.extract("Host: 127.0.0.1\n\rHost: 10.1.1.1\n\r")
+  #   # => ["127.0.0.1", "10.1.1.1"]
+  #
+  # @example Extract only IPv4 addresses from a large amount of text.
+  #   IPAddr.extract(text,:v4) do |ip|
+  #     puts ip
+  #   end
+  #   
   # @api public
   #
   def IPAddr.extract(text,version=nil,&block)
+    return enum_for(:extract,text,version).to_a unless block_given?
+
     regexp = case version
-             when :ipv4
-               IPV4_REGEXP
-             when :ipv6
-               IPV6_REGEXP
+             when :ipv4, :v4, 4
+               Regexp::IPv4
+             when :ipv6, :v6, 6
+               Regexp::IPv6
              else
-               REGEXP
+               Regexp::IP
              end
 
-    parser = StringScanner.new(text)
+    scanner = StringScanner.new(text)
 
-    if block_given?
-      yield parser.matched while parser.skip_until(regexp)
-      return nil
-    else
-      ips = []
-
-      ips << parser.matched while parser.skip_until(regexp)
-      return ips
+    while scanner.skip_until(regexp)
+      yield scanner.matched
     end
+
+    return nil
   end
 
   #
   # Iterates over each IP address within the IP Address range. Supports
   # both IPv4 and IPv6 address ranges.
@@ -118,32 +124,31 @@
     end
 
     return enum_for(:each,cidr_or_glob) unless block
 
     if cidr_or_glob.include?('::')
-      prefix = if cidr_or_glob[0,2] == '::'
+      separator = '::'
+      base      = 16
+
+      prefix = if cidr_or_glob.start_with?('::')
                  '::'
                else
                  ''
                end
 
-      separator = '::'
-      base = 16
-
       format = lambda { |address|
         prefix + address.map { |i| '%.2x' % i }.join('::')
       }
     else
       separator = '.'
-      base = 10
-
-      format = lambda { |address| address.join('.') }
+      base      = 10
+      format    = lambda { |address| address.join('.') }
     end
 
     # split the address
     segments = cidr_or_glob.split(separator)
-    ranges = []
+    ranges   = []
     
     # map the components of the address to numeric ranges
     segments.each do |segment|
       next if segment.empty?
 
@@ -157,27 +162,27 @@
                   segment.to_i(base)
                 end
     end
 
     # cycle through the address ranges
-    ranges.comprehension do |address|
-      yield format[address]
-    end
-
+    ranges.comprehension { |address| yield format[address] }
     return nil
   end
 
   #
   # Resolves the host-names for the IP address.
   #
+  # @param [String] nameserver
+  #   The optional nameserver to query.
+  #
   # @return [Array<String>]
   #   The host-names for the IP address.
   #
   # @api public
   #
-  def lookup
-    Resolv.getnames(self.to_s)
+  def lookup(nameserver=nil)
+    Resolv.resolver(nameserver).getnames(self.to_s)
   end
 
   #
   # Iterates over each IP address that is included in the addresses
   # netmask. Supports both IPv4 and IPv6 addresses.
@@ -199,15 +204,10 @@
   # @api public
   #
   def each
     return enum_for(:each) unless block_given?
 
-    case @family
-    when Socket::AF_INET
-      family_mask = IN4MASK
-    when Socket::AF_INET6
-      family_mask = IN6MASK
-    end
+    family_mask = MASKS[@family]
 
     (0..((~@mask_addr) & family_mask)).each do |i|
       yield _to_string(@addr | i)
     end
 