Implement configurable timeout option to Wayback Machine collector

lib/aquatone/collectors/wayback_machine.rb

@@ -6,18 +6,34 @@ class WaybackMachine < Aquatone::Collector
       self.meta = {
         :name         => "Wayback Machine",
         :author       => "Joel (@jolle)",
-        :description  => "Uses Wayback Machine by Internet Archive to find unique hostnames"
+        :description  => "Uses Wayback Machine by Internet Archive to find unique hostnames",
+        :cli_options  => {
+          "wayback-machine-timeout SECONDS" => "Timeout for Wayback Machine collector in seconds (default: 30)"
+        }
       }
 
-      def run
-        response = get_request("http://web.archive.org/cdx/search/cdx?url=*.#{url_escape(domain.name)}&output=json&fl=original&collapse=urlkey")
+      DEFAULT_TIMEOUT = 30.freeze
 
+      def run
+        response = nil
+        Timeout::timeout(timeout) do
+          response = get_request("http://web.archive.org/cdx/search/cdx?url=*.#{url_escape(domain.name)}&output=json&fl=original&collapse=urlkey")
+        end
         response.parsed_response.each do |page|
           if page[0] != "original"
               add_host(URI.parse(page[0]).host)
           end
         end
       end
+
+      private
+
+      def timeout
+        if has_cli_option?("wayback-machine-timeout")
+          return get_cli_option("wayback-machine-timeout").to_i
+        end
+        DEFAULT_TIMEOUT
+      end
     end
   end
 end

lib/aquatone/commands/discover.rb

@@ -84,6 +84,8 @@ def run_collectors
           rescue Aquatone::Collector::MissingKeyRequirement => e
             output(yellow("Skipped\n"))
             output(yellow(" -> #{e.message}\n"))
+          rescue Timeout::Error
+            output(red("Timed out\n"))
           rescue => e
             output(red("Error\n"))
             output(red(" -> #{e.message}\n"))