Implement Collector module defined CLI options

In order to make aquatone-discover and its collector modules more
flexible, collector modules can now define their own CLI options for
aquatone-discover which can be accessed within the modules. This has
made it possible to define the following new options:

 * --wordlist: Use a custom wordlist for dictionary collector instead of
   built-in wordlist
 * --censys-pages: Number of Censys API pages to process
 * --gtr-pages: Number of Google Transparency Report pages to process
 * --shodan-pages: Number of Shodan API pages to process
 * --netcraft-pages: Number of Netcraft pages to process

exe/aquatone-discover

@@ -121,6 +121,14 @@ OptionParser.new do |opts|
     options[:jitter] = v.to_f
   end
 
+  Aquatone::Collector.descendants.each do |collector|
+    collector.cli_options.each_pair do |option, description|
+      opts.on("--#{option}", description) do |v|
+        options[option.split(" ").first.gsub("-", "_").to_sym] = v
+      end
+    end
+  end
+
   opts.on("-h", "--help", "Show help") do
     puts opts
     exit 0

lib/aquatone/collector.rb

@@ -23,15 +23,21 @@ def self.descendants
       collectors.sort { |x, y| x.priority <=> y.priority }
     end
 
+    def self.cli_options
+      meta.key?(:cli_options) ? meta[:cli_options] : {}
+    end
+
     def self.sluggified_name
       return meta[:slug].downcase if meta[:slug]
       meta[:name].strip.downcase.gsub(/[^a-z0-9]+/, '-').gsub("--", "-")
     end
 
-    def initialize(domain)
+    def initialize(domain, options)
       check_key_requirements!
-      @domain = domain
-      @hosts  = []
+      @domain          = domain
+      @options         = options
+      @hosts           = []
+      @host_dictionary = {}
     end
 
     def run
@@ -52,7 +58,9 @@ def self.priority
     def add_host(host)
       host.downcase!
       return unless Aquatone::Validation.valid_domain_name?(host)
-      @hosts << host unless @hosts.include?(host)
+      return if @host_dictionary.key?(host)
+      @host_dictionary[host] = true
+      @hosts << host
     end
 
     def get_request(uri, options={})
@@ -83,6 +91,14 @@ def has_key?(name)
       Aquatone::KeyStore.key?(name)
     end
 
+    def get_cli_option(name)
+      @options[name.to_s.gsub("-", "_").to_sym]
+    end
+
+    def has_cli_option?(name)
+      @options.key?(name.to_s.gsub("-", "_").to_sym)
+    end
+
     def failure(message)
       fail Error, message
     end
@@ -101,6 +117,13 @@ def self.validate_metadata(meta)
       fail InvalidMetadataError, "Metadata is missing key: name" unless meta.key?(:name)
       fail InvalidMetadataError, "Metadata is missing key: author" unless meta.key?(:author)
       fail InvalidMetadataError, "Metadata is missing key: description" unless meta.key?(:description)
+      if meta.key?(:cli_options)
+        fail InvalidMetadataError, "Metadata CLI options is not a hash" unless meta[:cli_options].is_a?(Hash)
+        meta[:cli_options].each_pair do |option, description|
+          fail InvalidMetadataError, "CLI option name is not a string" unless option.is_a?(String)
+          fail InvalidMetadataError, "CLI option details is not a string" unless description.is_a?(String)
+        end
+      end
     end
   end
 end

lib/aquatone/collectors/censys.rb

@@ -6,11 +6,14 @@ class Censys < Aquatone::Collector
         :author       => "James McLean (@vortexau)",
         :description  => "Uses the Censys API to find hostnames in TLS certificates",
         :require_keys => ["censys_secret","censys_id"],
+        :cli_options  => {
+          "censys-pages PAGES" => "Number of Censys API pages to process (default: 10)"
+        }
       }
 
-      API_BASE_URI         = "https://www.censys.io/api/v1".freeze
-      API_RESULTS_PER_PAGE = 100.freeze
-      PAGES_TO_PROCESS     = 10.freeze
+      API_BASE_URI             = "https://www.censys.io/api/v1".freeze
+      API_RESULTS_PER_PAGE     = 100.freeze
+      DEFAULT_PAGES_TO_PROCESS = 10.freeze
 
       def run
         request_censys_page
@@ -21,10 +24,10 @@ def request_censys_page(page=1)
 
           # Censys expects Basic Auth for requests.
           auth = {
-              :username => get_key('censys_id'), 
+              :username => get_key('censys_id'),
               :password => get_key('censys_secret')
           }
-   
+
           # Define this is JSON content
           headers = {
             'Content-Type' => 'application/json',
@@ -41,11 +44,11 @@ def request_censys_page(page=1)
 
           # Search API documented at https://censys.io/api/v1/docs/search
           response = post_request(
-              "#{API_BASE_URI}/search/certificates", 
+              "#{API_BASE_URI}/search/certificates",
               query.to_json,
               {
                   :basic_auth => auth,
-                  :headers => headers 
+                  :headers => headers
               }
           )
 
@@ -75,9 +78,15 @@ def request_censys_page(page=1)
       end
 
       def next_page?(page, body)
-          page <= PAGES_TO_PROCESS && body["metadata"]["pages"] && API_RESULTS_PER_PAGE * page < body["metadata"]["count"].to_i
+          page <= pages_to_process && body["metadata"]["pages"] && API_RESULTS_PER_PAGE * page < body["metadata"]["count"].to_i
       end
 
+      def pages_to_process
+        if has_cli_option?("censys-pages")
+          return get_cli_option("censys-pages").to_i
+        end
+        DEFAULT_PAGES_TO_PROCESS
+      end
     end
   end
 end

lib/aquatone/collectors/dictionary.rb

@@ -4,13 +4,25 @@ class Dictionary < Aquatone::Collector
       self.meta = {
         :name        => "Dictionary",
         :author      => "Michael Henriksen (@michenriksen)",
-        :description => "Uses a dictionary to find hostnames"
+        :description => "Uses a dictionary to find hostnames",
+        :cli_options => {
+          "wordlist WORDLIST" => "OPTIONAL: wordlist/dictionary file to use for subdomain bruteforcing"
+        }
       }
 
-      DICTIONARY = File.join(Aquatone::AQUATONE_ROOT, "subdomains.lst").freeze
+      DEFAULT_DICTIONARY = File.join(Aquatone::AQUATONE_ROOT, "subdomains.lst").freeze
 
       def run
-        dictionary = File.open(DICTIONARY, "r")
+        if has_cli_option?("wordlist")
+          file = File.expand_path(get_cli_option("wordlist"))
+          if !File.readable?(file)
+            failure("Wordlist file #{file} is not readable or does not exist")
+          end
+          dictionary = File.open(file, "r")
+        else
+          dictionary = File.open(DEFAULT_DICTIONARY, "r")
+        end
+
         dictionary.each_line do |subdomain|
           add_host("#{subdomain.strip}.#{domain.name}")
         end

lib/aquatone/collectors/gtr.rb

@@ -5,15 +5,18 @@ class Gtr < Aquatone::Collector
         :name        => "Google Transparency Report",
         :author      => "Michael Henriksen (@michenriksen)",
         :description => "Uses Google Transparency Report to find hostnames",
-        :slug        => "gtr"
+        :slug        => "gtr",
+        :cli_options  => {
+          "gtr-pages PAGES" => "Number of Google Transparency Report pages to process (default: 30)"
+        }
       }
 
-      BASE_URI         = "https://www.google.com/transparencyreport/jsonp/ct/search"
-      PAGES_TO_PROCESS = 30.freeze
+      BASE_URI                 = "https://www.google.com/transparencyreport/jsonp/ct/search"
+      DEFAULT_PAGES_TO_PROCESS = 30.freeze
 
       def run
         token = nil
-        PAGES_TO_PROCESS.times do
+        pages_to_process.times do
           response = parse_response(request_page(token))
           response["results"].each do |result|
             host = result["subject"]
@@ -53,6 +56,13 @@ def valid_host?(host)
         return false unless host.end_with?(".#{domain.name}")
         true
       end
+
+      def pages_to_process
+        if has_cli_option?("gtr-pages")
+          return get_cli_option("gtr-pages").to_i
+        end
+        DEFAULT_PAGES_TO_PROCESS
+      end
     end
   end
 end

lib/aquatone/collectors/netcraft.rb

@@ -4,18 +4,21 @@ class Netcraft < Aquatone::Collector
       self.meta = {
         :name        => "Netcraft",
         :author      => "Michael Henriksen (@michenriksen)",
-        :description => "Uses searchdns.netcraft.com to find hostnames"
+        :description => "Uses searchdns.netcraft.com to find hostnames",
+        :cli_options  => {
+          "netcraft-pages PAGES" => "Number of Netcraft pages to process (default: 10)"
+        }
       }
 
-      BASE_URI         = "http://searchdns.netcraft.com/".freeze
-      HOSTNAME_REGEX   = /<a href="http:\/\/(.*?)\/" rel="nofollow">/.freeze
-      RESULTS_PER_PAGE = 20.freeze
-      PAGES_TO_PROCESS = 10.freeze
+      BASE_URI                 = "http://searchdns.netcraft.com/".freeze
+      HOSTNAME_REGEX           = /<a href="http:\/\/(.*?)\/" rel="nofollow">/.freeze
+      RESULTS_PER_PAGE         = 20.freeze
+      DEFAULT_PAGES_TO_PROCESS = 10.freeze
 
       def run
         last  = nil
         count = 0
-        PAGES_TO_PROCESS.times do |i|
+        pages_to_process.times do |i|
           page = i + 1
           if page == 1
             uri = "#{BASE_URI}/?restriction=site+contains&host=*.#{url_escape(domain.name)}&lookup=wait..&position=limited"
@@ -43,6 +46,13 @@ def extract_hostnames_from_response(body)
         end
         hosts
       end
+
+      def pages_to_process
+        if has_cli_option?("netcraft-pages")
+          return get_cli_option("netcraft-pages").to_i
+        end
+        DEFAULT_PAGES_TO_PROCESS
+      end
     end
   end
 end

lib/aquatone/collectors/shodan.rb

@@ -5,12 +5,15 @@ class Shodan < Aquatone::Collector
         :name         => "Shodan",
         :author       => "Michael Henriksen (@michenriksen)",
         :description  => "Uses the Shodan API to find hostnames",
-        :require_keys => ["shodan"]
+        :require_keys => ["shodan"],
+        :cli_options  => {
+          "shodan-pages PAGES" => "Number of Shodan API pages to process (default: 10)"
+        }
       }
 
-      API_BASE_URI         = "https://api.shodan.io/shodan".freeze
-      API_RESULTS_PER_PAGE = 100.freeze
-      PAGES_TO_PROCESS     = 10.freeze
+      API_BASE_URI             = "https://api.shodan.io/shodan".freeze
+      API_RESULTS_PER_PAGE     = 100.freeze
+      DEFAULT_PAGES_TO_PROCESS = 10.freeze
 
       def run
         request_shodan_page
@@ -38,7 +41,14 @@ def construct_uri(query, page)
       end
 
       def next_page?(page, body)
-        page <= PAGES_TO_PROCESS && body["total"] && API_RESULTS_PER_PAGE * page < body["total"].to_i
+        page <= pages_to_process && body["total"] && API_RESULTS_PER_PAGE * page < body["total"].to_i
+      end
+
+      def pages_to_process
+        if has_cli_option?("shodan-pages")
+          return get_cli_option("shodan-pages").to_i
+        end
+        DEFAULT_PAGES_TO_PROCESS
       end
     end
   end

lib/aquatone/commands/discover.rb

@@ -77,7 +77,7 @@ def run_collectors
           next if skip_collector?(collector)
           output("Running collector: #{bold(collector.meta[:name])}... ")
           begin
-            collector_instance = collector.new(@domain)
+            collector_instance = collector.new(@domain, options)
             hosts = collector_instance.execute!
             output("Done (#{hosts.count} #{hosts.count == 1 ? 'host' : 'hosts'})\n")
             @hosts += hosts