lib/ronin/network/dns.rb in ronin-support-0.4.1 vs lib/ronin/network/dns.rb in ronin-support-0.5.0.rc1

@@ -40,11 +40,11 @@
       end
 
       #
       # Sets the DNS nameserver to be queried.
       #
-      # @param [IPAddr, String, nil]
+      # @param [IPAddr, String, nil] address
       #   The address of the nameserver.
       #
       # @return [String, nil]
       #   The address of the new nameserver.
       #
@@ -82,22 +82,32 @@
       #   The hostname to lookup.
       #
       # @param [String, nil] nameserver
       #   Optional DNS nameserver to query.
       #
+      # @yield [address]
+      #   If a block is given and the hostname was resolved, the address will
+      #   be passed to the block.
+      #
+      # @yieldparam [String] address
+      #   The address of the hostname.
+      #
       # @return [String, nil]
       #   The address of the hostname.
       #
       # @api public
       #
       def dns_lookup(hostname,nameserver=DNS.nameserver)
-        resolv = dns_resolver(nameserver)
+        hostname = hostname.to_s
+        resolv   = dns_resolver(nameserver)
+        address  = begin
+                     resolv.getaddress(hostname).to_s
+                   rescue Resolv::ResolvError
+                   end
 
-        begin
-          resolv.getaddress(hostname.to_s).to_s
-        rescue Resolv::ResolvError
-        end
+        yield(address) if (block_given? && address)
+        return address
       end
 
       #
       # Looks up all addresses of a hostname.
       #
@@ -105,17 +115,28 @@
       #   The hostname to lookup.
       #
       # @param [String, nil] nameserver
       #   Optional DNS nameserver to query.
       #
+      # @yield [address]
+      #   If a block is given, each resolved address will be passed to the
+      #   block.
+      #
+      # @yieldparam [String] address
+      #   A address of the hostname.
+      #
       # @return [Array<String>]
       #   The addresses of the hostname.
       #
       # @api public
       #
-      def dns_lookup_all(hostname,nameserver=DNS.nameserver)
-        dns_resolver(nameserver).getaddresses(hostname.to_s).map(&:to_s)
+      def dns_lookup_all(hostname,nameserver=DNS.nameserver,&block)
+        hostname  = hostname.to_s
+        addresses = dns_resolver(nameserver).getaddresses(hostname).map(&:to_s)
+
+        addresses.each(&block) if block
+        return addresses
       end
 
       #
       # Looks up the hostname of the address.
       #
@@ -123,22 +144,32 @@
       #   The address to lookup.
       #
       # @param [String, nil] nameserver
       #   Optional DNS nameserver to query.
       #
+      # @yield [hostname]
+      #   If a block is given and a hostname was found for the address,
+      #   the resolved hostname will be passed to the block.
+      #
+      # @yieldparam [String] hostname
+      #   The hostname of the address.
+      #
       # @return [String, nil]
       #   The hostname of the address.
       #
       # @api public
       #
       def dns_reverse_lookup(address,nameserver=DNS.nameserver)
-        resolv = dns_resolver(nameserver)
+        address  = address.to_s
+        resolv   = dns_resolver(nameserver)
+        hostname = begin
+                     resolv.getname(address).to_s
+                   rescue Resolv::ResolvError
+                   end
 
-        begin
-          resolv.getname(address.to_s).to_s
-        rescue Resolv::ResolvError
-        end
+        yield(hostname) if (block_given? && hostname)
+        return hostname
       end
 
       #
       # Looks up all hostnames associated with the address.
       #
@@ -146,16 +177,27 @@
       #   The address to lookup.
       #
       # @param [String, nil] nameserver
       #   Optional DNS nameserver to query.
       #
+      # @yield [hostname]
+      #   If a block is given and hostnames were found for the address,
+      #   each hostname will be passed to the block.
+      #
+      # @yieldparam [String] hostname
+      #   A hostname of the address.
+      #
       # @return [Array<String>]
       #   The hostnames of the address.
       #
       # @api public
       #
-      def dns_reverse_lookup_all(address,nameserver=DNS.nameserver)
-        dns_resolver(nameserver).getnames(address.to_s).map(&:to_s)
+      def dns_reverse_lookup_all(address,nameserver=DNS.nameserver,&block)
+        address   = address.to_s
+        hostnames = dns_resolver(nameserver).getnames(address).map(&:to_s)
+
+        hostnames.each(&block) if block
+        return hostnames
       end
     end
   end
 end